From 52b1259240803590895324a1f9c723ac5f4112b0 Mon Sep 17 00:00:00 2001 From: talnufai Date: Thu, 20 Aug 2026 10:43:13 +0530 Subject: [PATCH 1/4] fix: capture correct chunked inputs and all iterations in write_io_files - Pass chunk_inputs instead of full padded tensor during prefill - Use indexed subdirs (prefill_{i}, decode_{n}) to preserve all steps - Remove self._write_io_dir = None that suppressed subsequent captures Signed-off-by: talnufai --- .../generation/text_generation_inference.py | 14 ++++++---- .../transformers/models/modeling_auto.py | 27 ++++++++++++------- 2 files changed, 26 insertions(+), 15 deletions(-) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 17c992064c..d6e828a63a 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -848,7 +848,9 @@ def run_prefill(self, prompt, generation_len, prefill_logit_bs=1, decode_batch_i outputs = self._session.run(chunk_inputs) if self._write_io_dir is not None: - write_io_files(inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + write_io_files( + chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False + ) return ( outputs, position_ids, @@ -1020,8 +1022,9 @@ def run_decode( outputs = self._session.run(decode_inputs) if self._write_io_dir is not None: - write_io_files(decode_inputs, outputs, self._write_io_dir, "decode", "aic_batch_io", True, False) - self._write_io_dir = None + write_io_files( + decode_inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False + ) # Prepare inputs for next iteration decode_inputs["input_ids"] = self._fetch_next_token_id(outputs) @@ -1055,8 +1058,9 @@ def generate_decode_stream(self, decode_inputs, generation_len, automation): outputs = self._session.run(decode_inputs) if self._write_io_dir is not None: - write_io_files(decode_inputs, outputs, self._write_io_dir, "decode", "aic_batch_io", True, False) - self._write_io_dir = None + write_io_files( + decode_inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False + ) # Prepare inputs for next iteration decode_inputs["input_ids"] = outputs["logits"].argmax(2) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 213986d41e..87e3a994bb 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -2565,7 +2565,9 @@ def kv_offload_generate( chunk_inputs["image_idx"] = outputs["image_idx_output"] if self._write_io_dir is not None: - write_io_files(lang_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + write_io_files( + chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False + ) prefill_time = perf_counter() - lang_start + vision_end - vision_start # Skip inputs/outputs again @@ -2624,14 +2626,15 @@ def kv_offload_generate( outputs = lang_session.run(lang_inputs) if self._write_io_dir is not None: - write_io_files(lang_inputs, outputs, self._write_io_dir, "decode", "aic_batch_io", True, False) - self._write_io_dir = None + write_io_files( + lang_inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False + ) # Prepare inputs for next iteration lang_inputs["input_ids"] = outputs["logits"].argmax(2) lang_inputs["position_ids"] += 1 if "mm_token_type_ids" in lang_inputs: - lang_inputs["mm_token_type_ids"] = np.zeros_like( + lang_inputs["mm_token_ids"] = np.zeros_like( lang_inputs["input_ids"], dtype=lang_inputs["mm_token_type_ids"].dtype ) generated_ids[:, num_token] = lang_inputs["input_ids"].squeeze(1) @@ -3181,7 +3184,9 @@ def cloud_ai_100_generate( outputs = qpc_session.run(chunk_inputs) if self._write_io_dir is not None: - write_io_files(chunk_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + write_io_files( + chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False + ) chunk_inputs["image_idx"] = outputs["image_idx_output"] @@ -3226,8 +3231,9 @@ def cloud_ai_100_generate( outputs = qpc_session.run(inputs) if self._write_io_dir is not None: - write_io_files(inputs, outputs, self._write_io_dir, "decode", "aic_batch_io", True, False) - self._write_io_dir = None + write_io_files( + inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False + ) # Prepare inputs for next iteration inputs["input_ids"] = outputs["logits"].argmax(2) @@ -5195,7 +5201,7 @@ def generate( outputs = self.qpc_session.run(inputs) if self._write_io_dir is not None: - write_io_files(inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + write_io_files(inputs, outputs, self._write_io_dir, "prefill_0", "aic_batch_io", True, False) # array to hold generated tokens generated_ids = np.full((self.batch_size, generation_len + 1), self.model.config.eos_token_id) @@ -5213,8 +5219,9 @@ def generate( for num_tokens in range(generation_len): outputs = self.qpc_session.run(inputs) if self._write_io_dir is not None: - write_io_files(inputs, outputs, self._write_io_dir, "decode", "aic_batch_io", True, False) - self._write_io_dir = None + write_io_files( + inputs, outputs, self._write_io_dir, f"decode_{num_tokens}", "aic_batch_io", True, False + ) logits = outputs["logits"] next_token = logits.argmax(-1) From 6959ec2190adc5c862336513d9ddaf0a4c3a0797 Mon Sep 17 00:00:00 2001 From: talnufai Date: Thu, 20 Aug 2026 11:24:55 +0530 Subject: [PATCH 2/4] fix(write_io_files): capture vision encoder I/O and fix subdir naming collision Two write_io_files gaps in the VLM inference paths: 1. modeling_auto.py (kv_offload_generate): vision_session.run() inputs and outputs were never captured. Added write_io_files call immediately after the vision encoder run using subdir "vision_prefill_0". 2. vlm_generation.py (run_prefill_multi_frame_specialization): each per-frame vision_session.run() was never captured. Added write_io_files call inside the frame loop using subdir f"vision_prefill_{i}". Naming: vision encoder subdirs are prefixed "vision_prefill_" to avoid collision with the lang chunked prefill loop which writes to "prefill_{i}". Using the same "prefill_{i}" namespace for both sessions would cause the lang chunk to overwrite the vision .raw files on disk while producing two JSON entries pointing to the same paths, silently corrupting the captured trace. Signed-off-by: talnufai --- QEfficient/generation/vlm_generation.py | 6 +++++- QEfficient/transformers/models/modeling_auto.py | 4 ++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index a56187a6b6..aa07148986 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -397,7 +397,7 @@ def _execute_chunked_prefill( chunk_image_idx = outputs["image_idx_output"] if self._write_io_dir is not None: - write_io_files(lang_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + write_io_files(chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False) # Prepare decode-time cross_attention_mask if "cross_attention_mask" in lang_inputs: @@ -673,6 +673,10 @@ def run_prefill_multi_frame_specialization( for i in range(num_frames): chunk_inputs["pixel_values"] = vision_inputs["pixel_values"][i * vision_size : (i + 1) * vision_size] chunk_outputs = self._vision_session.run(chunk_inputs) + if self._write_io_dir is not None: + write_io_files( + chunk_inputs, chunk_outputs, self._write_io_dir, f"vision_frame_{i}", "aic_batch_io", True, False + ) if i == 0: vision_outputs = chunk_outputs else: diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 87e3a994bb..b09708769e 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -2462,6 +2462,10 @@ def kv_offload_generate( vision_outputs = {} if vision_inputs: vision_outputs = vision_session.run(vision_inputs) + if self._write_io_dir is not None: + write_io_files( + vision_inputs, vision_outputs, self._write_io_dir, "vision_prefill_0", "aic_batch_io", True, False + ) vision_end = perf_counter() lang_inputs = {k: v for k, v in inputs.items() if k not in vision_inputs} From d0b1b0eb318279253b73f6888ce9d02de3a6b8a0 Mon Sep 17 00:00:00 2001 From: talnufai Date: Thu, 20 Aug 2026 17:56:12 +0530 Subject: [PATCH 3/4] fix(write_io): timestamp io_dir per generate() call and fix dual-QPC onnx_path indexing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add get_io_dir() helper that appends an EST timestamp to io_dir so each run writes to its own subdirectory. Replace all _write_io_dir assignments with it. Also fix _QEffAutoModelForImageTextToTextDualQPC using self.onnx_path[1] — replace with explicit self.lang_model.onnx_path. Signed-off-by: talnufai --- .../transformers/models/modeling_auto.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index b09708769e..ed272ba155 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -11,6 +11,7 @@ from pathlib import Path from time import perf_counter from typing import List, Optional, Union +from datetime import datetime, timezone, timedelta import numpy as np import onnx @@ -109,6 +110,15 @@ torch.float32: np.float32, } +def get_io_dir(onnx_path: str) -> str: + """Return a timestamped io_dir path under the model's onnx directory. + + Format: /io_dir/ + """ + est = timezone(timedelta(hours=-5)) + timestamp = datetime.now(est).strftime("%Y%m%d_%H%M%S") + return os.path.join(os.path.dirname(onnx_path), "io_dir", timestamp) + def _resolve_torch_dtype(kwargs: dict) -> None: """ @@ -731,7 +741,7 @@ def generate( torch.Tensor or np.ndarray Output from the AI 100 or PyTorch runtime. The type depends on the runtime and model. """ - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.onnx_path) if write_io else None # AI_100 runtime if runtime_ai100: @@ -2292,7 +2302,7 @@ def generate( raise NotImplementedError("PyTorch execution is not supported yet for this model!") write_io = kwargs.pop("write_io", False) - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path[1]), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.lang_model.onnx_path) if write_io else None # Use VisionLanguageGeneration for image-prompt pairs if (processor and images) or (tokenizer and prompts) or multi_specs or num_frames: @@ -3062,7 +3072,7 @@ def generate( if not runtime_ai100: raise NotImplementedError("PyTorch execution is not supported yet for this model!") - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.onnx_path) if write_io else None return self.cloud_ai_100_generate( inputs=inputs, device_ids=device_ids, generation_len=generation_len, streamer=streamer @@ -4816,7 +4826,7 @@ def generate( If `runtime_ai100` is False. """ write_io = kwargs.pop("write_io", False) - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.onnx_path) if write_io else None if runtime_ai100: if not isinstance(self.qpc_path, Path): @@ -5173,7 +5183,7 @@ def generate( if not isinstance(self.qpc_path, Path): raise TypeError("Please run compile API first!") - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.onnx_path) if write_io else None inputs = self.auto_correct_inputs(inputs) if self.qpc_session is None: @@ -5475,7 +5485,7 @@ def generate( Returns: :dict: Output from the ``AI_100`` or ``PyTorch`` runtime. """ - self._write_io_dir = os.path.join(os.path.dirname(self.onnx_path), "io_dir") if write_io else None + self._write_io_dir = get_io_dir(self.onnx_path) if write_io else None # AI_100 runtime if runtime_ai100: From 3a6f30ab604e63b7d4770a9c6eda48dc70de7d7a Mon Sep 17 00:00:00 2001 From: talnufai Date: Tue, 8 Sep 2026 05:01:42 -0700 Subject: [PATCH 4/4] Lint & Format Signed-off-by: talnufai --- .../generation/text_generation_inference.py | 4 +--- QEfficient/generation/vlm_generation.py | 8 +++++++- .../transformers/models/modeling_auto.py | 19 ++++++------------- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index d6e828a63a..b76b156124 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -848,9 +848,7 @@ def run_prefill(self, prompt, generation_len, prefill_logit_bs=1, decode_batch_i outputs = self._session.run(chunk_inputs) if self._write_io_dir is not None: - write_io_files( - chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False - ) + write_io_files(chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False) return ( outputs, position_ids, diff --git a/QEfficient/generation/vlm_generation.py b/QEfficient/generation/vlm_generation.py index aa07148986..5fdff52400 100755 --- a/QEfficient/generation/vlm_generation.py +++ b/QEfficient/generation/vlm_generation.py @@ -675,7 +675,13 @@ def run_prefill_multi_frame_specialization( chunk_outputs = self._vision_session.run(chunk_inputs) if self._write_io_dir is not None: write_io_files( - chunk_inputs, chunk_outputs, self._write_io_dir, f"vision_frame_{i}", "aic_batch_io", True, False + chunk_inputs, + chunk_outputs, + self._write_io_dir, + f"vision_frame_{i}", + "aic_batch_io", + True, + False, ) if i == 0: vision_outputs = chunk_outputs diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index ed272ba155..9db5c7a7a4 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -8,10 +8,10 @@ import math import os import warnings +from datetime import datetime, timedelta, timezone from pathlib import Path from time import perf_counter from typing import List, Optional, Union -from datetime import datetime, timezone, timedelta import numpy as np import onnx @@ -110,6 +110,7 @@ torch.float32: np.float32, } + def get_io_dir(onnx_path: str) -> str: """Return a timestamped io_dir path under the model's onnx directory. @@ -2579,9 +2580,7 @@ def kv_offload_generate( chunk_inputs["image_idx"] = outputs["image_idx_output"] if self._write_io_dir is not None: - write_io_files( - chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False - ) + write_io_files(chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False) prefill_time = perf_counter() - lang_start + vision_end - vision_start # Skip inputs/outputs again @@ -3198,9 +3197,7 @@ def cloud_ai_100_generate( outputs = qpc_session.run(chunk_inputs) if self._write_io_dir is not None: - write_io_files( - chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False - ) + write_io_files(chunk_inputs, outputs, self._write_io_dir, f"prefill_{i}", "aic_batch_io", True, False) chunk_inputs["image_idx"] = outputs["image_idx_output"] @@ -3245,9 +3242,7 @@ def cloud_ai_100_generate( outputs = qpc_session.run(inputs) if self._write_io_dir is not None: - write_io_files( - inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False - ) + write_io_files(inputs, outputs, self._write_io_dir, f"decode_{num_token}", "aic_batch_io", True, False) # Prepare inputs for next iteration inputs["input_ids"] = outputs["logits"].argmax(2) @@ -5233,9 +5228,7 @@ def generate( for num_tokens in range(generation_len): outputs = self.qpc_session.run(inputs) if self._write_io_dir is not None: - write_io_files( - inputs, outputs, self._write_io_dir, f"decode_{num_tokens}", "aic_batch_io", True, False - ) + write_io_files(inputs, outputs, self._write_io_dir, f"decode_{num_tokens}", "aic_batch_io", True, False) logits = outputs["logits"] next_token = logits.argmax(-1)