diff --git a/profile_streaming_bottleneck_gpu7.py b/profile_streaming_bottleneck_gpu7.py index 9711844..c472f9d 100644 --- a/profile_streaming_bottleneck_gpu7.py +++ b/profile_streaming_bottleneck_gpu7.py @@ -205,6 +205,32 @@ def _patch_bound_method( setattr(owner, name, types.MethodType(wrapped, owner)) +def _patch_component_call(component: Any, stage: str) -> None: + """Time ``component(...)`` despite Python's class-level special lookup.""" + + component_class = type(component) + original = component_class.__call__ + + @functools.wraps(original) + def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any: + if instance is component: + return PROFILER.measure(stage, original, instance, *args, **kwargs) + return original(instance, *args, **kwargs) + + # Qwen's processor, feature extractor, video processor, and tokenizer use + # distinct classes in Transformers 4.52.3. Guard anyway so an unexpected + # shared class cannot silently replace an earlier component timer. + marker = "_omni_bottleneck_profile_stage" + previous_stage = getattr(component_class, marker, None) + if previous_stage is not None: + raise RuntimeError( + f"Cannot separately profile {stage}: {component_class.__name__} " + f"is already patched as {previous_stage}" + ) + setattr(component_class, marker, stage) + setattr(component_class, "__call__", wrapped) + + def _install_profiling_hooks() -> None: # One-time input setup. _patch_function(direct, "_decode_audio_exact", "audio_decode") @@ -215,6 +241,24 @@ def _install_profiling_hooks() -> None: _patch_function(direct.CausalSourceAV, "frames", "media_frames") _patch_function(direct.CausalSourceAV, "audio_slice", "media_audio_slice") _patch_function(core, "_prepare_conversation", "processor") + _patch_function(core, "process_mm_info", "processor_mm_info") + _patch_function( + core, + "_processor_batch_convert_to_tensors", + "processor_tensorize", + ) + _patch_function( + core, + "_processor_batch_to_device", + "processor_h2d", + cuda_sync=True, + ) + _patch_function( + core, + "_processor_batch_to_dtype", + "processor_dtype_cast", + cuda_sync=True, + ) # Per-state GPU/cache work. Leaf GPU calls synchronize for attribution. _patch_function(core, "clone_cache", "cache_clone", cuda_sync=True) @@ -283,6 +327,40 @@ def _install_profiling_hooks() -> None: _patch_bound_method( model.model, "forward", "thinker_forward", cuda_sync=True ) + processor_components = { + "processor_qwen_call": processor, + "processor_audio": getattr(processor, "feature_extractor", None), + "processor_video": getattr(processor, "video_processor", None), + "processor_tokenizer": getattr(processor, "tokenizer", None), + } + missing_components = [ + stage + for stage, component in processor_components.items() + if component is None + ] + if missing_components: + raise RuntimeError( + "Installed Qwen processor is missing profiling components: " + + ", ".join(missing_components) + ) + for stage, component in processor_components.items(): + _patch_component_call(component, stage) + _patch_bound_method( + processor, + "apply_chat_template", + "processor_chat_template", + cuda_sync=False, + ) + _patch_bound_method( + processor, + "replace_multimodal_special_tokens", + "processor_placeholder_expansion", + cuda_sync=False, + ) + PROFILER.metadata["processor_component_classes"] = { + stage: type(component).__name__ + for stage, component in processor_components.items() + } torch.cuda.reset_peak_memory_stats() return model, processor @@ -382,7 +460,8 @@ def _print_summary( print("\n=== Qwen Omni bottleneck profile (inclusive timings) ===") print( "Nested rows intentionally overlap; compare leaf rows such as " - "processor/audio_tower/video_tower/thinker_forward/cache_clone." + "processor_audio/processor_video/processor_tokenizer/processor_h2d " + "and the model-tower/Thinker rows." ) print( f"overall={overall_seconds:.3f}s extract={extract_seconds:.3f}s " @@ -420,7 +499,15 @@ def _print_summary( leaf_stages = ( "media_frames", - "processor", + "processor_mm_info", + "processor_chat_template", + "processor_audio", + "processor_video", + "processor_placeholder_expansion", + "processor_tokenizer", + "processor_tensorize", + "processor_h2d", + "processor_dtype_cast", "audio_tower", "video_tower", "tmrope", @@ -433,7 +520,15 @@ def _print_summary( bottleneck = max(totals, key=totals.get) advice = { "media_frames": "bulk-decode/cache the selected 2-FPS resized frames", - "processor": "prefetch or microbatch the ten causal prefixes per 2-s block", + "processor_mm_info": "cache the frame-list/media packaging metadata", + "processor_chat_template": "cache the invariant rendered chat template", + "processor_audio": "precompute or incrementally slice exact audio features", + "processor_video": "cache transformed 2-FPS frames or batch equal grids", + "processor_placeholder_expansion": "cache token-layout templates by prefix shape", + "processor_tokenizer": "cache token IDs by multimodal prefix shape", + "processor_tensorize": "return PyTorch tensors directly from the processor", + "processor_h2d": "use pinned buffers and non-blocking host-to-device copies", + "processor_dtype_cast": "produce the target floating dtype before transfer", "audio_tower": "microbatch the ten 0.2-s prefixes in each 2-s block", "video_tower": "cache unique 2-FPS visual inputs and batch equal shapes", "tmrope": "cache position IDs for the ten repeated prefix shapes", diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 7d6dff7..65a7938 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -565,12 +565,31 @@ def _prepare_conversation( ): processor_inputs["images"] = images batch = processor(**processor_inputs) - batch = batch.convert_to_tensors("pt") + batch = _processor_batch_convert_to_tensors(batch) # BatchFeature.to(dtype) changes floating tensors only; IDs stay integral. - batch = batch.to(model.device).to(model.dtype) + batch = _processor_batch_to_device(batch, model.device) + batch = _processor_batch_to_dtype(batch, model.dtype) return dict(batch) +def _processor_batch_convert_to_tensors(batch: Any) -> Any: + """Materialize the processor's NumPy/list outputs as PyTorch tensors.""" + + return batch.convert_to_tensors("pt") + + +def _processor_batch_to_device(batch: Any, device: torch.device) -> Any: + """Move one processor batch to the model device.""" + + return batch.to(device) + + +def _processor_batch_to_dtype(batch: Any, dtype: torch.dtype) -> Any: + """Cast only floating processor tensors to the model dtype.""" + + return batch.to(dtype) + + def _readout_query_ids( processor: Qwen2_5OmniProcessor, device: torch.device,