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 5, 02:15 UTC
[Model] Use mm_features for Keye-VL and Keye-1.5-VL M-RoPE (#39869)
Signed-off-by: Lalit Laxminarayan Bangad <lalitbangad@gmail.com>
lalit10 committed
Apr 16, 2026 at 02:16 UTC
17d87168d27c2d2a99c544b57ea046629bcdc48f
4 files changed
+425
-158
tests/model_executor/test_keye_mrope.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass, field
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.keye import KeyeForConditionalGeneration
9
+from vllm.multimodal.inputs import (
10
+ MultiModalFeatureSpec,
11
+ MultiModalFieldElem,
12
+ MultiModalKwargsItem,
13
+ PlaceholderRange,
14
+)
15
+
16
+pytestmark = pytest.mark.skip_global_cleanup
17
+
18
+
19
+@pytest.fixture(autouse=True, scope="module")
20
+def _force_cpu_default_device():
21
+ original = torch.get_default_device()
22
+ torch.set_default_device("cpu")
23
+ yield
24
+ torch.set_default_device(original)
25
+
26
+
27
+@dataclass
28
+class DummyVisionConfig:
29
+ spatial_merge_size: int = 2
30
+
31
+
32
+@dataclass
33
+class DummyConfig:
34
+ vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
35
+
36
+
37
+def make_model(config: DummyConfig) -> KeyeForConditionalGeneration:
38
+ model = object.__new__(KeyeForConditionalGeneration)
39
+ model.config = config
40
+ return model
41
+
42
+
43
+def make_mm_feature(
44
+ *,
45
+ modality: str,
46
+ offset: int,
47
+ length: int,
48
+ grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
49
+) -> MultiModalFeatureSpec:
50
+ field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
51
+ return MultiModalFeatureSpec(
52
+ data=MultiModalKwargsItem(
53
+ {
54
+ field_name: MultiModalFieldElem(
55
+ data=torch.tensor(grid_thw),
56
+ field=None, # HACK.
57
+ ),
58
+ }
59
+ ),
60
+ modality=modality,
61
+ identifier="DUMMY",
62
+ mm_position=PlaceholderRange(offset=offset, length=length),
63
+ )
64
+
65
+
66
+def test_get_mrope_input_positions_text_only():
67
+ model = make_model(DummyConfig())
68
+
69
+ positions, delta = model.get_mrope_input_positions(
70
+ input_tokens=[11, 12, 13, 14, 15],
71
+ mm_features=[],
72
+ )
73
+
74
+ expected = torch.tensor(
75
+ [
76
+ [0, 1, 2, 3, 4],
77
+ [0, 1, 2, 3, 4],
78
+ [0, 1, 2, 3, 4],
79
+ ]
80
+ )
81
+
82
+ assert torch.equal(positions, expected)
83
+ assert delta == 0
84
+
85
+
86
+def test_get_mrope_input_positions_single_image():
87
+ model = make_model(DummyConfig())
88
+ mm_features = [
89
+ make_mm_feature(
90
+ modality="image",
91
+ offset=1,
92
+ length=4,
93
+ grid_thw=(1, 4, 4),
94
+ )
95
+ ]
96
+
97
+ positions, delta = model.get_mrope_input_positions(
98
+ input_tokens=[10, 20, 21, 22, 23, 30, 31],
99
+ mm_features=mm_features,
100
+ )
101
+
102
+ expected = torch.tensor(
103
+ [
104
+ [0, 1, 1, 1, 1, 3, 4],
105
+ [0, 1, 1, 2, 2, 3, 4],
106
+ [0, 1, 2, 1, 2, 3, 4],
107
+ ]
108
+ )
109
+
110
+ assert torch.equal(positions, expected)
111
+ assert delta == -2
112
+
113
+
114
+def test_get_mrope_input_positions_interleaved_image_and_video():
115
+ model = make_model(DummyConfig())
116
+ mm_features = [
117
+ make_mm_feature(
118
+ modality="image",
119
+ offset=1,
120
+ length=4,
121
+ grid_thw=(1, 4, 4),
122
+ ),
123
+ make_mm_feature(
124
+ modality="video",
125
+ offset=7,
126
+ length=4,
127
+ grid_thw=[(2, 4, 2)],
128
+ ),
129
+ ]
130
+
131
+ positions, delta = model.get_mrope_input_positions(
132
+ input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51],
133
+ mm_features=mm_features,
134
+ )
135
+
136
+ expected = torch.tensor(
137
+ [
138
+ [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10],
139
+ [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10],
140
+ [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10],
141
+ ]
142
+ )
143
+
144
+ assert torch.equal(positions, expected)
145
+ assert delta == -2
tests/model_executor/test_keye_vl1_5_mrope.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass, field
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration
9
+from vllm.multimodal.inputs import (
10
+ MultiModalFeatureSpec,
11
+ MultiModalFieldElem,
12
+ MultiModalKwargsItem,
13
+ PlaceholderRange,
14
+)
15
+
16
+pytestmark = pytest.mark.skip_global_cleanup
17
+
18
+
19
+@pytest.fixture(autouse=True, scope="module")
20
+def _force_cpu_default_device():
21
+ original = torch.get_default_device()
22
+ torch.set_default_device("cpu")
23
+ yield
24
+ torch.set_default_device(original)
25
+
26
+
27
+@dataclass
28
+class DummyVisionConfig:
29
+ spatial_merge_size: int = 2
30
+
31
+
32
+@dataclass
33
+class DummyConfig:
34
+ vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
35
+
36
+
37
+def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration:
38
+ model = object.__new__(KeyeVL1_5ForConditionalGeneration)
39
+ model.config = config
40
+ return model
41
+
42
+
43
+def make_mm_feature(
44
+ *,
45
+ modality: str,
46
+ offset: int,
47
+ length: int,
48
+ grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
49
+ is_embed: list[bool] | None = None,
50
+) -> MultiModalFeatureSpec:
51
+ field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
52
+ return MultiModalFeatureSpec(
53
+ data=MultiModalKwargsItem(
54
+ {
55
+ field_name: MultiModalFieldElem(
56
+ data=torch.tensor(grid_thw),
57
+ field=None, # HACK.
58
+ ),
59
+ }
60
+ ),
61
+ modality=modality,
62
+ identifier="DUMMY",
63
+ mm_position=PlaceholderRange(
64
+ offset=offset,
65
+ length=length,
66
+ is_embed=None if is_embed is None else torch.tensor(is_embed),
67
+ ),
68
+ )
69
+
70
+
71
+def test_get_mrope_input_positions_text_only():
72
+ model = make_model(DummyConfig())
73
+
74
+ positions, delta = model.get_mrope_input_positions(
75
+ input_tokens=[11, 12, 13, 14, 15],
76
+ mm_features=[],
77
+ )
78
+
79
+ expected = torch.tensor(
80
+ [
81
+ [0, 1, 2, 3, 4],
82
+ [0, 1, 2, 3, 4],
83
+ [0, 1, 2, 3, 4],
84
+ ]
85
+ )
86
+
87
+ assert torch.equal(positions, expected)
88
+ assert delta == 0
89
+
90
+
91
+def test_get_mrope_input_positions_single_image():
92
+ model = make_model(DummyConfig())
93
+ mm_features = [
94
+ make_mm_feature(
95
+ modality="image",
96
+ offset=1,
97
+ length=4,
98
+ grid_thw=(1, 4, 4),
99
+ )
100
+ ]
101
+
102
+ positions, delta = model.get_mrope_input_positions(
103
+ input_tokens=[10, 20, 21, 22, 23, 30, 31],
104
+ mm_features=mm_features,
105
+ )
106
+
107
+ expected = torch.tensor(
108
+ [
109
+ [0, 1, 1, 1, 1, 3, 4],
110
+ [0, 1, 1, 2, 2, 3, 4],
111
+ [0, 1, 2, 1, 2, 3, 4],
112
+ ]
113
+ )
114
+
115
+ assert torch.equal(positions, expected)
116
+ assert delta == -2
117
+
118
+
119
+def test_get_mrope_input_positions_video_uses_embed_ranges():
120
+ model = make_model(DummyConfig())
121
+ mm_features = [
122
+ make_mm_feature(
123
+ modality="video",
124
+ offset=1,
125
+ length=8,
126
+ grid_thw=[(2, 4, 2)],
127
+ is_embed=[False, False, True, True, False, False, True, True],
128
+ )
129
+ ]
130
+
131
+ positions, delta = model.get_mrope_input_positions(
132
+ input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41],
133
+ mm_features=mm_features,
134
+ )
135
+
136
+ expected = torch.tensor(
137
+ [
138
+ [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
139
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
140
+ [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
141
+ ]
142
+ )
143
+
144
+ assert torch.equal(positions, expected)
145
+ assert delta == 0
vllm/model_executor/models/keye.py
+80
-79
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from abc import abstractmethod
5
-from collections.abc import Iterable, Mapping, Sequence
5
+from collections.abc import Iterable, Iterator, Mapping, Sequence
6
from functools import partial
7
from typing import Annotated, Any, Literal, TypeAlias, TypeVar
8
@@ -1595,91 +1595,92 @@ class KeyeForConditionalGeneration(
1595
self._process_video_embeds(video_type, video_grid_thw, pixel_values_videos)
1596
)
1597
1598
+ @staticmethod
1599
+ def _split_video_grid_thw(
1600
+ grid_thw: torch.Tensor | list[list[int]] | list[int],
1601
+ ) -> list[list[int]]:
1602
+ """
1603
+ Split video grid_thw along the t dimension into per-frame rows.
1604
+
1605
+ This preserves Keye's current M-RoPE behavior, where a video is emitted
1606
+ as consecutive frame-level multimodal blocks rather than a single block
1607
+ spanning the whole video.
1608
+ """
1609
+ if isinstance(grid_thw, list):
1610
+ if len(grid_thw) == 0:
1611
+ return []
1612
+ if isinstance(grid_thw[0], int):
1613
+ grid_thw = torch.tensor([grid_thw], dtype=torch.long)
1614
+ else:
1615
+ grid_thw = torch.tensor(grid_thw, dtype=torch.long)
1616
+ elif grid_thw.ndim == 1:
1617
+ grid_thw = grid_thw.unsqueeze(0)
1618
+
1619
+ if grid_thw.numel() == 0:
1620
+ return []
1621
+
1622
+ t, hw = grid_thw[:, 0], grid_thw[:, 1:]
1623
+ ones = torch.ones_like(hw[:, :1])
1624
+ out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0)
1625
+ return out.tolist()
1626
+
1627
+ def iter_mm_grid_thw(
1628
+ self, mm_features: list[MultiModalFeatureSpec]
1629
+ ) -> Iterator[tuple[int, int, int, int]]:
1630
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
1631
+
1632
+ for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
1633
+ if mm_feature.data is None:
1634
+ raise ValueError("M-RoPE calculation requires multimodal feature data")
1635
+
1636
+ if mm_feature.modality == "image":
1637
+ grid_thw = mm_feature.data["image_grid_thw"].data
1638
+ if isinstance(grid_thw, torch.Tensor):
1639
+ if grid_thw.ndim == 2:
1640
+ assert grid_thw.shape[0] == 1
1641
+ t, h, w = grid_thw[0].tolist()
1642
+ else:
1643
+ t, h, w = grid_thw.tolist()
1644
+ else:
1645
+ if isinstance(grid_thw[0], list):
1646
+ assert len(grid_thw) == 1
1647
+ t, h, w = grid_thw[0]
1648
+ else:
1649
+ t, h, w = grid_thw
1650
+
1651
+ yield (
1652
+ mm_feature.mm_position.offset,
1653
+ t,
1654
+ h // spatial_merge_size,
1655
+ w // spatial_merge_size,
1656
+ )
1657
+ elif mm_feature.modality == "video":
1658
+ current_offset = mm_feature.mm_position.offset
1659
+ for t, h, w in self._split_video_grid_thw(
1660
+ mm_feature.data["video_grid_thw"].data
1661
+ ):
1662
+ llm_grid_h = h // spatial_merge_size
1663
+ llm_grid_w = w // spatial_merge_size
1664
+ yield (current_offset, t, llm_grid_h, llm_grid_w)
1665
+ current_offset += t * llm_grid_h * llm_grid_w
1666
+ else:
1667
+ raise ValueError(f"Unsupported modality: {mm_feature.modality}")
1668
+
1669
def get_mrope_input_positions(
1670
self,
1671
input_tokens: list[int],
1672
mm_features: list[MultiModalFeatureSpec],
1673
) -> tuple[torch.Tensor, int]:
1603
- kwargs = MultiModalFeatureSpec.gather_kwargs(
1604
- mm_features,
1605
- {"image_grid_thw", "video_grid_thw"},
1606
- )
1607
- image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])]
1608
- video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])]
1609
-
1610
- if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0:
1611
- video_grid_thw = video_grid_thw[0]
1612
-
1613
- def split_thw(grid_thw: torch.Tensor | list[int]) -> list[list[int]]:
1614
- """
1615
- Split grid_thw along the t dimension.
1616
-
1617
- Args:
1618
- grid_thw: shape [N, 3] tensor or nested list of [t, h, w].
1619
-
1620
- Returns:
1621
- List of [1, h, w] rows, repeated t times for each original row.
1622
- """
1623
-
1624
- if isinstance(grid_thw, list):
1625
- grid_thw = torch.tensor(grid_thw, dtype=torch.long)
1626
-
1627
- if grid_thw.numel() == 0:
1628
- return []
1629
-
1630
- t, hw = grid_thw[:, 0], grid_thw[:, 1:]
1631
- ones = torch.ones_like(hw[:, :1]) # [N,1]
1632
- out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0)
1633
- return out.tolist()
1634
-
1635
- video_grid_thw = split_thw(video_grid_thw)
1636
-
1637
- hf_config = self.config
1638
- image_token_id = hf_config.image_token_id
1639
- video_token_id = hf_config.video_token_id
1640
- spatial_merge_size = hf_config.vision_config.spatial_merge_size
1641
-
1642
- image_nums = len(image_grid_thw)
1643
- frame_nums = len(video_grid_thw)
1674
llm_pos_ids_list: list = []
1645
-
1675
st = 0
1647
- remain_images, remain_frames = image_nums, frame_nums
1648
-
1649
- image_index, video_index = 0, 0
1650
- for _ in range(image_nums + frame_nums):
1651
- if remain_images > 0:
1652
- try:
1653
- ed_image = input_tokens.index(image_token_id, st)
1654
- except ValueError:
1655
- ed_image = len(input_tokens) + 1
1656
- else:
1657
- ed_image = len(input_tokens) + 1
1658
- if remain_frames > 0:
1659
- try:
1660
- ed_video = input_tokens.index(video_token_id, st)
1661
- except ValueError:
1662
- ed_video = len(input_tokens) + 1
1663
- else:
1664
- ed_video = len(input_tokens) + 1
1676
1666
- if ed_image < ed_video:
1667
- t, h, w = image_grid_thw[image_index]
1668
- image_index += 1
1669
- remain_images -= 1
1670
- ed = ed_image
1671
- else:
1672
- t, h, w = video_grid_thw[video_index]
1673
- video_index += 1
1674
- remain_frames -= 1
1675
- ed = ed_video
1676
-
1677
- llm_grid_t, llm_grid_h, llm_grid_w = (
1678
- t,
1679
- h // spatial_merge_size,
1680
- w // spatial_merge_size,
1681
- )
1682
- text_len = ed - st
1677
+ for (
1678
+ offset,
1679
+ llm_grid_t,
1680
+ llm_grid_h,
1681
+ llm_grid_w,
1682
+ ) in self.iter_mm_grid_thw(mm_features):
1683
+ text_len = offset - st
1684
1685
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1686
llm_pos_ids_list.append(
@@ -1711,7 +1712,7 @@ class KeyeForConditionalGeneration(
1712
llm_pos_ids_list.append(
1713
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
1714
)
1714
- st = ed + llm_grid_t * llm_grid_h * llm_grid_w
1715
+ st = offset + llm_grid_t * llm_grid_h * llm_grid_w
1716
1717
if st < len(input_tokens):
1718
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
vllm/model_executor/models/keye_vl1_5.py
+55
-79
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import itertools
4
-from collections.abc import Mapping, Sequence
4
+from collections.abc import Iterator, Mapping, Sequence
5
from functools import partial
6
from typing import Annotated, Any, Literal, TypeAlias
7
@@ -608,91 +608,67 @@ class KeyeVL1_5ForConditionalGeneration(
608
new_video_embeds.append(video_embeds[start:end])
609
return tuple(new_video_embeds)
610
611
+ def iter_mm_grid_thw(
612
+ self, mm_features: list[MultiModalFeatureSpec]
613
+ ) -> Iterator[tuple[int, int, int, int]]:
614
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
615
+
616
+ for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
617
+ if mm_feature.data is None:
618
+ raise ValueError("M-RoPE calculation requires multimodal feature data")
619
+
620
+ embed_ranges = mm_feature.mm_position.extract_embeds_range()
621
+ if mm_feature.modality == "image":
622
+ assert len(embed_ranges) == 1
623
+ grid_thw = mm_feature.data["image_grid_thw"].data
624
+ if isinstance(grid_thw, torch.Tensor):
625
+ if grid_thw.ndim == 2:
626
+ assert grid_thw.shape[0] == 1
627
+ t, h, w = grid_thw[0].tolist()
628
+ else:
629
+ t, h, w = grid_thw.tolist()
630
+ else:
631
+ if isinstance(grid_thw[0], list):
632
+ assert len(grid_thw) == 1
633
+ t, h, w = grid_thw[0]
634
+ else:
635
+ t, h, w = grid_thw
636
+
637
+ yield (
638
+ embed_ranges[0][0],
639
+ t,
640
+ h // spatial_merge_size,
641
+ w // spatial_merge_size,
642
+ )
643
+ elif mm_feature.modality == "video":
644
+ split_video_grids = split_thw(mm_feature.data["video_grid_thw"].data)
645
+ assert len(embed_ranges) == split_video_grids.shape[0]
646
+ for (start_idx, end_idx), (t, h, w) in zip(
647
+ embed_ranges, split_video_grids.tolist()
648
+ ):
649
+ llm_grid_h = h // spatial_merge_size
650
+ llm_grid_w = w // spatial_merge_size
651
+ num_mm_tokens = t * llm_grid_h * llm_grid_w
652
+ assert end_idx - start_idx + 1 == num_mm_tokens
653
+ yield (start_idx, t, llm_grid_h, llm_grid_w)
654
+ else:
655
+ raise ValueError(f"Unsupported modality: {mm_feature.modality}")
656
+
657
def get_mrope_input_positions(
658
self,
659
input_tokens: list[int],
660
mm_features: list[MultiModalFeatureSpec],
661
) -> tuple[torch.Tensor, int]:
616
- kwargs = MultiModalFeatureSpec.gather_kwargs(
617
- mm_features,
618
- {"image_grid_thw", "video_grid_thw"},
619
- )
620
- image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])]
621
- video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])]
622
-
623
- if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0:
624
- video_grid_thw = video_grid_thw[0]
625
-
626
- def split_thw(grid_thw: torch.Tensor | list[int]) -> list[list[int]]:
627
- """
628
- Split grid_thw along the t dimension.
629
-
630
- Args:
631
- grid_thw: shape [N, 3] tensor or nested list of [t, h, w].
632
-
633
- Returns:
634
- List of [1, h, w] rows, repeated t times for each original row.
635
- """
636
-
637
- if isinstance(grid_thw, list):
638
- grid_thw = torch.tensor(grid_thw, dtype=torch.long)
639
-
640
- if grid_thw.numel() == 0:
641
- return []
642
-
643
- t, hw = grid_thw[:, 0], grid_thw[:, 1:]
644
- ones = torch.ones_like(hw[:, :1]) # [N,1]
645
- out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0)
646
- return out.tolist()
647
-
648
- video_grid_thw = split_thw(video_grid_thw)
649
-
650
- hf_config = self.config
651
- image_token_id = hf_config.image_token_id
652
- video_token_id = hf_config.video_token_id
653
- spatial_merge_size = hf_config.vision_config.spatial_merge_size
654
-
655
- image_nums = len(image_grid_thw)
656
- frame_nums = len(video_grid_thw)
662
llm_pos_ids_list: list = []
658
-
663
st = 0
660
- remain_images, remain_frames = image_nums, frame_nums
661
-
662
- image_index, video_index = 0, 0
663
- for _ in range(image_nums + frame_nums):
664
- if remain_images > 0:
665
- try:
666
- ed_image = input_tokens.index(image_token_id, st)
667
- except ValueError:
668
- ed_image = len(input_tokens) + 1
669
- else:
670
- ed_image = len(input_tokens) + 1
671
- if remain_frames > 0:
672
- try:
673
- ed_video = input_tokens.index(video_token_id, st)
674
- except ValueError:
675
- ed_video = len(input_tokens) + 1
676
- else:
677
- ed_video = len(input_tokens) + 1
664
679
- if ed_image < ed_video:
680
- t, h, w = image_grid_thw[image_index]
681
- image_index += 1
682
- remain_images -= 1
683
- ed = ed_image
684
- else:
685
- t, h, w = video_grid_thw[video_index]
686
- video_index += 1
687
- remain_frames -= 1
688
- ed = ed_video
689
-
690
- llm_grid_t, llm_grid_h, llm_grid_w = (
691
- t,
692
- h // spatial_merge_size,
693
- w // spatial_merge_size,
694
- )
695
- text_len = ed - st
665
+ for (
666
+ offset,
667
+ llm_grid_t,
668
+ llm_grid_h,
669
+ llm_grid_w,
670
+ ) in self.iter_mm_grid_thw(mm_features):
671
+ text_len = offset - st
672
673
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
674
llm_pos_ids_list.append(
@@ -724,7 +700,7 @@ class KeyeVL1_5ForConditionalGeneration(
700
llm_pos_ids_list.append(
701
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
702
)
727
- st = ed + llm_grid_t * llm_grid_h * llm_grid_w
703
+ st = offset + llm_grid_t * llm_grid_h * llm_grid_w
704
705
if st < len(input_tokens):
706
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0