Revert "[V1] Scatter and gather placeholders in the model runner" (#16075)

This commit is contained in:
Roger Wang
2025-04-04 14:50:57 -07:00
committed by GitHub
parent f5722a5052
commit af51d80fa1
42 changed files with 942 additions and 496 deletions

View File

@@ -108,46 +108,16 @@ class PromptUpdateDetails(Generic[_S]):
full: _S
"""The full content."""
is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]] = None
features: _S
"""
Given :attr:`full`, return a boolean mask of shape `(len(full),)`
indicating which positions of `full` to assign embeddings to.
`None` (default) means to assign embeddings to all positions of `full`.
The embeddings are obtained by calling
:class:`SupportsMultiModal.get_multimodal_embeddings`.
The part of the content that corresponds to feature placeholders;
this will be replaced by the output of the vision encoder during model
inference.
"""
@staticmethod
def from_seq(seq: _S) -> "PromptUpdateDetails[_S]":
return PromptUpdateDetails(full=seq)
@staticmethod
def select_text(
seq: _S,
embed_text: str,
) -> "PromptUpdateDetails[_S]":
def is_embed(full: "_BoundPromptSequence") -> torch.Tensor:
embed_token_ids = encode_tokens(full.tokenizer, embed_text)
return torch.isin(
torch.tensor(full.token_ids),
torch.tensor(embed_token_ids),
)
return PromptUpdateDetails(full=seq, is_embed=is_embed)
@staticmethod
def select_token_id(
seq: _S,
embed_token_id: int,
) -> "PromptUpdateDetails[_S]":
return PromptUpdateDetails(
full=seq,
is_embed=lambda f: torch.tensor(f.token_ids) == embed_token_id,
)
return PromptUpdateDetails(full=seq, features=seq)
PromptUpdateInfo = Union[PromptSeq, PromptUpdateDetails]
@@ -436,7 +406,7 @@ class _BoundPromptSequence:
@dataclass
class _BoundPromptContent:
full: _BoundPromptSequence
is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]]
features: _BoundPromptSequence
@dataclass
@@ -496,8 +466,10 @@ class BoundPromptUpdate:
bound_full = _BoundPromptSequence.from_seq(self.tokenizer,
content.full)
bound_features = _BoundPromptSequence.from_seq(self.tokenizer,
content.features)
bound_content = _BoundPromptContent(full=bound_full,
is_embed=content.is_embed)
features=bound_features)
if cache_key is not None:
self._content_cache[cache_key] = bound_content
@@ -633,19 +605,15 @@ class PlaceholderFeaturesInfo:
item_idx: int
start_idx: int
tokens: list[int]
is_embed: Optional[torch.Tensor]
@property
def length(self) -> int:
return len(self.tokens)
def to_range(self) -> PlaceholderRange:
# TODO: Is it worth it to optimize this by stripping the
# leading and ending positions where `is_embed=False`?
return PlaceholderRange(
offset=self.start_idx,
length=self.length,
is_embed=self.is_embed,
)
@@ -838,17 +806,22 @@ def _iter_placeholders(
continue
if prompt[start_idx:end_idx_full] == content_tokens_full:
content_is_embed = content.is_embed
if content_is_embed is not None:
content_is_embed = content_is_embed(content.full)
content_tokens_feat = content.features.token_ids
yield PlaceholderFeaturesInfo(
modality=modality,
item_idx=item_idx,
start_idx=start_idx,
tokens=content_tokens_full,
is_embed=content_is_embed,
)
try:
match = next(
iter_token_matches(content_tokens_full,
content_tokens_feat))
yield PlaceholderFeaturesInfo(
modality=modality,
item_idx=item_idx,
start_idx=start_idx + match.start_idx,
tokens=content_tokens_feat,
)
except StopIteration:
raise AssertionError(
f"{content_tokens_feat=} should be a "
f"subsequence of {content_tokens_full=}") from None
# Exclude overlapping matches
start_idx = end_idx_full