[MM] Keep more metadata tensors on CPU (#52827)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Nick Hill
2026-08-18 19:56:33 -07:00
committed by GitHub
parent aabc1a0a0b
commit deeeae75d0
38 changed files with 148 additions and 83 deletions
+1 -1
View File
@@ -268,7 +268,7 @@ class Cohere2VisionMultiModalProcessor(
num_patches = hf_inputs.get("num_patches", torch.empty(0))
return dict(
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
num_patches=MultiModalFieldConfig.batched("image"),
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_embeds=MultiModalFieldConfig.batched("image"),
)
+3 -1
View File
@@ -320,7 +320,9 @@ class DeepseekOCRMultiModalProcessor(
patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
images_spatial_crop=MultiModalFieldConfig.batched("image"),
images_spatial_crop=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
images_crop=MultiModalFieldConfig.flat_from_sizes(
"image", patches_per_image
),
+3 -1
View File
@@ -191,7 +191,9 @@ class DeepseekOCR2MultiModalProcessor(
patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
images_spatial_crop=MultiModalFieldConfig.batched("image"),
images_spatial_crop=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
images_crop=MultiModalFieldConfig.flat_from_sizes(
"image", patches_per_image
),
+3 -1
View File
@@ -270,7 +270,9 @@ class DeepseekVL2MultiModalProcessor(
return dict(
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
images_spatial_crop=MultiModalFieldConfig.batched("image"),
images_spatial_crop=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
image_embeds=MultiModalFieldConfig.batched("image"),
)
+2 -2
View File
@@ -1246,11 +1246,11 @@ class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessin
pixel_values=MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes
),
image_grid_thw=MultiModalFieldConfig.batched("image"),
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_grid_sizes
),
video_grid_thw=MultiModalFieldConfig.batched("video"),
video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
)
+1 -1
View File
@@ -260,7 +260,7 @@ class FireRedASR2MultiModalProcessor(
return dict(
input_features=MultiModalFieldConfig.batched("audio"),
speech_lengths=MultiModalFieldConfig.batched("audio"),
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
def _get_prompt_updates(
+2 -2
View File
@@ -482,8 +482,8 @@ class FireRedLIDMultiModalProcessor(
) -> Mapping[str, MultiModalFieldConfig]:
return dict(
input_features=MultiModalFieldConfig.batched("audio"),
speech_lengths=MultiModalFieldConfig.batched("audio"),
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
def _get_prompt_updates(
+1 -1
View File
@@ -779,7 +779,7 @@ class FunASRMultiModalProcessor(BaseMultiModalProcessor[FunASRProcessingInfo]):
return dict(
input_features=MultiModalFieldConfig.batched("audio"),
speech_lengths=MultiModalFieldConfig.batched("audio"),
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
def _get_prompt_updates(
+3 -1
View File
@@ -695,7 +695,9 @@ class FunAudioChatMultiModalProcessor(
) -> Mapping[str, MultiModalFieldConfig]:
return {
"speech_ids": MultiModalFieldConfig.batched("audio"),
"speech_attention_mask": MultiModalFieldConfig.batched("audio"),
"speech_attention_mask": MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
),
"input_features": MultiModalFieldConfig.batched("audio"),
"feature_attention_mask": MultiModalFieldConfig.batched("audio"),
"feature_exist_mask": MultiModalFieldConfig.batched("audio"),
+2 -2
View File
@@ -625,13 +625,13 @@ def _glmasr_field_config(
feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(
"audio", chunk_counts, dim=0
),
chunk_counts=MultiModalFieldConfig.batched("audio"),
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
return dict(
audio_embeds=MultiModalFieldConfig.batched("audio"),
input_features=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
chunk_counts=MultiModalFieldConfig.batched("audio"),
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
@@ -398,7 +398,7 @@ class Granite4VisionMultiModalProcessor(
) -> Mapping[str, MultiModalFieldConfig]:
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
image_sizes=MultiModalFieldConfig.batched("image"),
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
@@ -324,10 +324,14 @@ class HCXVisionMultiModalProcessor(BaseMultiModalProcessor[HCXVisionProcessingIn
) -> Mapping[str, MultiModalFieldConfig]:
fields = dict(
pixel_values_images=MultiModalFieldConfig.batched("image"),
image_sizes_images=MultiModalFieldConfig.batched("image"),
vision_query_lengths_images=MultiModalFieldConfig.batched("image"),
image_sizes_images=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
vision_query_lengths_images=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
pixel_values_videos=MultiModalFieldConfig.batched("video"),
vision_query_lengths_videos=MultiModalFieldConfig.batched("video"),
vision_query_lengths_videos=MultiModalFieldConfig.batched(
"video", keep_on_cpu=True
),
)
return fields
+1 -1
View File
@@ -440,7 +440,7 @@ class IsaacMultiModalProcessor(BaseMultiModalProcessor):
"pixel_values": MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes
),
"image_grid_thw": MultiModalFieldConfig.batched("image"),
"image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
}
def _get_prompt_updates(
+2 -2
View File
@@ -561,8 +561,8 @@ class KananaVMultiModalProcessor(BaseMultiModalProcessor[KananaVProcessingInfo])
mm_fields_config = dict(
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", pixel_sizes),
vision_grid_thw=MultiModalFieldConfig.batched("image"),
image_token_thw=MultiModalFieldConfig.batched("image"),
vision_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_token_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
return mm_fields_config
+5 -3
View File
@@ -309,13 +309,15 @@ def _keye_field_config(
return dict(
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
image_grid_thw=MultiModalFieldConfig.batched("image"),
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_patches
),
video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_num_patches),
video_grid_thw=MultiModalFieldConfig.flat_from_sizes("video", video_num_grids),
num_frames=MultiModalFieldConfig.batched("video"),
video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_grids, keep_on_cpu=True
),
num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
)
+1 -1
View File
@@ -201,7 +201,7 @@ class KimiAudioDummyInputsBuilder(BaseDummyInputsBuilder[KimiAudioProcessingInfo
# Field config for Kimi-Audio multimodal data
_KIMIAUDIO_FIELD_CONFIG = {
"whisper_input_features": MultiModalFieldConfig.batched("audio"),
"feature_attention_mask": MultiModalFieldConfig.batched("audio"),
"feature_attention_mask": MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
}
+1 -1
View File
@@ -251,7 +251,7 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
pixel_values=MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes
),
image_grid_hws=MultiModalFieldConfig.batched("image"),
image_grid_hws=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
def _get_prompt_updates(
@@ -431,20 +431,20 @@ def _create_field_factory(
image_embeds=MultiModalFieldConfig.flat_from_sizes(
"image", image_embed_grid_sizes
),
image_grid_thw=MultiModalFieldConfig.batched("image"),
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
# OV2 first-class MM kwarg: per-patch (t,h,w)
# positions required by the 3-D vision RoPE.
patch_positions=MultiModalFieldConfig.flat_from_sizes(
"image", image_pixel_grid_sizes
"image", image_pixel_grid_sizes, keep_on_cpu=True
),
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_patch_sizes
),
video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_frames
"video", video_num_frames, keep_on_cpu=True
),
patch_positions_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_patch_sizes
"video", video_patch_sizes, keep_on_cpu=True
),
video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
+9 -5
View File
@@ -893,17 +893,21 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing
merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
fields: dict[str, MultiModalFieldConfig] = dict(
**_create_qwen2vl_field_factory(merge_size)(hf_inputs),
second_per_grid_ts=MultiModalFieldConfig.batched("video"),
video_start_times=MultiModalFieldConfig.batched("video"),
second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_start_times=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
audio_features=MultiModalFieldConfig.batched("audio"),
audio_token_lens=MultiModalFieldConfig.batched("audio"),
audio_token_lens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
# video_audio fields: only present when video_audio content was processed
if "video_audio_n_segs" in hf_inputs:
fields["video_audio_n_segs"] = MultiModalFieldConfig.batched("video")
fields["video_audio_n_segs"] = MultiModalFieldConfig.batched(
"video", keep_on_cpu=True
)
# video_audio_seg_lens: list of per-video 1D tensors, batched("video")
if "video_audio_seg_lens" in hf_inputs:
fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched("video")
fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched(
"video", keep_on_cpu=True
)
if "va_audio_features" in hf_inputs:
fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio")
return fields
+4 -4
View File
@@ -472,12 +472,12 @@ def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:
def _minicpmv_field_config(hf_inputs: Mapping[str, torch.Tensor]):
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
image_sizes=MultiModalFieldConfig.batched("image"),
tgt_sizes=MultiModalFieldConfig.batched("image"),
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_embeds=MultiModalFieldConfig.batched("image"),
video_pixel_values=MultiModalFieldConfig.batched("video"),
video_image_sizes=MultiModalFieldConfig.batched("video"),
video_tgt_sizes=MultiModalFieldConfig.batched("video"),
video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_embeds=MultiModalFieldConfig.batched("video"),
)
+6 -4
View File
@@ -75,15 +75,17 @@ from .vision import is_vit_use_data_parallel
def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]):
fields = dict(
pixel_values=MultiModalFieldConfig.batched("image"),
tgt_sizes=MultiModalFieldConfig.batched("image"),
tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_embeds=MultiModalFieldConfig.batched("image"),
video_pixel_values=MultiModalFieldConfig.batched("video"),
video_image_sizes=MultiModalFieldConfig.batched("video"),
video_tgt_sizes=MultiModalFieldConfig.batched("video"),
video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_embeds=MultiModalFieldConfig.batched("video"),
)
if "use_vit_merger" in hf_inputs:
fields["use_vit_merger"] = MultiModalFieldConfig.batched("image")
fields["use_vit_merger"] = MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
)
return fields
+2 -2
View File
@@ -653,8 +653,8 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo])
pixel_values=MultiModalFieldConfig.flat_from_sizes(
"image", patches_per_image
),
patches_per_image=MultiModalFieldConfig.batched("image"),
aspect_ratios=MultiModalFieldConfig.batched("image"),
patches_per_image=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
aspect_ratios=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
def _get_prompt_updates(
+4 -2
View File
@@ -1208,8 +1208,10 @@ class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):
images=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
image_masks=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
image_input_idx=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
num_crops=MultiModalFieldConfig.batched("image"),
img_patch_id=MultiModalFieldConfig.shared("image", num_images),
num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
img_patch_id=MultiModalFieldConfig.shared(
"image", num_images, keep_on_cpu=True
),
)
def _get_prompt_updates(
+12 -8
View File
@@ -2127,26 +2127,30 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):
image_token_pooling=MultiModalFieldConfig.flat_from_sizes(
"image", image_num_pooled_patches
),
image_num_crops=MultiModalFieldConfig.batched("image"),
image_num_pooled_patches=MultiModalFieldConfig.batched("image"),
image_num_patches=MultiModalFieldConfig.batched("image"),
image_num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_num_pooled_patches=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_tokens=MultiModalFieldConfig.flat_from_sizes(
"image", num_image_tokens
),
num_image_tokens=MultiModalFieldConfig.batched("image"),
num_image_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_crops
),
video_token_pooling=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_pooled_patches
),
video_num_crops=MultiModalFieldConfig.batched("video"),
video_num_pooled_patches=MultiModalFieldConfig.batched("video"),
video_num_patches=MultiModalFieldConfig.batched("video"),
video_num_crops=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_num_pooled_patches=MultiModalFieldConfig.batched(
"video", keep_on_cpu=True
),
video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
video_tokens=MultiModalFieldConfig.flat_from_sizes(
"video", num_video_tokens
),
num_video_tokens=MultiModalFieldConfig.batched("video"),
num_video_tokens=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
)
def _get_prompt_updates(
@@ -352,9 +352,12 @@ def _mtd_field_config(
audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
"audio",
audio_chunk_counts,
keep_on_cpu=True,
),
audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
audio_token_lengths=MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
),
audio_chunk_counts=MultiModalFieldConfig.batched("audio"),
audio_token_lengths=MultiModalFieldConfig.batched("audio"),
)
return fields
+6 -2
View File
@@ -378,12 +378,16 @@ class MuseGlimmerMultiModalProcessor(
if "image_pixel_values" in hf_inputs:
fields.update(
image_pixel_values=MultiModalFieldConfig.batched("image"),
image_feature_sizes=MultiModalFieldConfig.batched("image"),
image_feature_sizes=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
)
if "video_pixel_values" in hf_inputs:
fields.update(
video_pixel_values=MultiModalFieldConfig.batched("video"),
video_feature_sizes=MultiModalFieldConfig.batched("video"),
video_feature_sizes=MultiModalFieldConfig.batched(
"video", keep_on_cpu=True
),
)
return fields
@@ -386,10 +386,12 @@ class NanoNemotronVLMultiModalProcessor(
return dict(
pixel_values_flat=pixel_values_flat,
image_num_patches=MultiModalFieldConfig.batched("image"),
image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
image_embeds=MultiModalFieldConfig.batched("image"),
num_tokens_per_image=MultiModalFieldConfig.batched("image"),
imgs_sizes=MultiModalFieldConfig.batched("image"),
num_tokens_per_image=MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
),
imgs_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
def _get_video_fields_config(self, hf_inputs: BatchFeature):
@@ -399,9 +401,9 @@ class NanoNemotronVLMultiModalProcessor(
pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes(
"video", video_num_patches
),
video_num_patches=MultiModalFieldConfig.batched("video"),
frames_indices=MultiModalFieldConfig.batched("video"),
frame_duration_ms=MultiModalFieldConfig.batched("video"),
video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
frames_indices=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
frame_duration_ms=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
)
def _get_audio_fields_config(self, hf_inputs: BatchFeature):
+1 -1
View File
@@ -386,7 +386,7 @@ class OvisMultiModalProcessor(BaseMultiModalProcessor[OvisProcessingInfo]):
) -> Mapping[str, MultiModalFieldConfig]:
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
grids=MultiModalFieldConfig.batched("image"),
grids=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
indicator_tokens=MultiModalFieldConfig.batched("image"),
)
+1 -1
View File
@@ -279,7 +279,7 @@ class PaddleOCRVLMultiModalProcessor(
) -> Mapping[str, MultiModalFieldConfig]:
return dict(
pixel_values=MultiModalFieldConfig.batched("image"),
image_grid_thw=MultiModalFieldConfig.batched("image"),
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
)
def _get_prompt_updates(
+1 -1
View File
@@ -918,7 +918,7 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):
input_image_embeds=MultiModalFieldConfig.batched("image"),
image_attention_mask=MultiModalFieldConfig.batched("image"),
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
num_img_tokens=MultiModalFieldConfig.batched("image"),
num_img_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
input_audio_embeds=MultiModalFieldConfig.batched("audio"),
)
@@ -301,8 +301,12 @@ def create_qwen2_5_omni_thinker_field_factory(
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
"audio", audio_feature_lengths, dim=1
),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
),
audio_feature_lengths=MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
),
pixel_values=MultiModalFieldConfig.flat_from_sizes(
"image", image_pixel_grid_sizes
),
@@ -318,7 +322,9 @@ def create_qwen2_5_omni_thinker_field_factory(
),
video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos),
use_audio_in_video=MultiModalFieldConfig.shared(
"video", num_videos, keep_on_cpu=True
),
)
return _qwen2_5_omni_thinker_field_config
@@ -1031,7 +1037,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
audio_input: Qwen2_5OmniAudioFeatureInputs,
) -> torch.Tensor:
input_features = audio_input["input_features"]
audio_feature_lengths = audio_input["audio_feature_lengths"]
# audio_feature_lengths is keep_on_cpu; the audio tower derives
# device placement from feature_lens, so move it explicitly.
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
input_features.device, non_blocking=True
)
audio_feat_lengths, audio_output_lengths = (
self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
+1 -1
View File
@@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
return dict(
audio_embeds=MultiModalFieldConfig.batched("audio"),
input_features=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
+7 -3
View File
@@ -254,8 +254,8 @@ def _qwen3asr_field_config(hf_inputs: Mapping[str, torch.Tensor]):
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
"audio", audio_feature_lengths, dim=1
),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
audio_feature_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
@@ -459,7 +459,11 @@ class Qwen3ASRForConditionalGeneration(
audio_input: Qwen2_5OmniAudioFeatureInputs,
) -> torch.Tensor:
input_features = audio_input["input_features"]
audio_feature_lengths = audio_input["audio_feature_lengths"]
# audio_feature_lengths is keep_on_cpu; the audio tower derives
# device placement from feature_lens, so move it explicitly.
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
input_features.device, non_blocking=True
)
audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
@@ -1553,7 +1553,11 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
audio_input: Qwen2_5OmniAudioFeatureInputs,
) -> tuple[torch.Tensor, ...]:
input_features = audio_input["input_features"]
audio_feature_lengths = audio_input["audio_feature_lengths"]
# audio_feature_lengths is keep_on_cpu; the audio tower derives
# device placement from feature_lens, so move it explicitly.
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
input_features.device, non_blocking=True
)
audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
+1 -1
View File
@@ -202,7 +202,7 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo])
),
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
patch_newline_mask=MultiModalFieldConfig.flat_from_sizes(
"image", num_patches
"image", num_patches, keep_on_cpu=True
),
)
@@ -300,11 +300,17 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
# Keep these as batched, as they always have batch size as first dim
if "audio" in modalities:
mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched("audio")
mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
)
if "image" in modalities:
mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image")
mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
)
# TODO: route to "video" once the video modality is supported
mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image")
mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
)
mm_fields["num_image_patches"] = MultiModalFieldConfig.batched(
"image", keep_on_cpu=True
)
@@ -755,6 +761,12 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
num_image_patches = kwargs.pop("num_image_patches")
# grid_thw fields are registered keep_on_cpu; restore the on-device
# placement that HF get_image_features implementations expect.
for key, value in kwargs.items():
if isinstance(value, torch.Tensor):
kwargs[key] = value.to(pixel_values.device, non_blocking=True)
# The underlying HuggingFace `get_image_features` implementations
# contain model-internal syncs (e.g. Idefics3 filters all-zero
# padding images via boolean-mask indexing, LlavaOnevision
+1 -1
View File
@@ -251,7 +251,7 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
"audio", num_chunks, keep_on_cpu=True
),
# num_chunks can convert audio_chunked to audio batch dimension
audio_num_chunks=MultiModalFieldConfig.batched("audio"),
audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
audio_embeds=MultiModalFieldConfig.batched("audio"),
)
+2 -2
View File
@@ -307,12 +307,12 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo])
return dict(
# Ragged per-image patches, grouped by num_patches.
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
num_patches=MultiModalFieldConfig.batched("image"),
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
# Ragged per-audio frames, grouped by num_audio_tokens.
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
"audio", num_audio_tokens
),
num_audio_tokens=MultiModalFieldConfig.batched("audio"),
num_audio_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
def _get_prompt_updates(