Skip to content

fix(write_io): capture complete generate() I/O traces for LLM and VLM models #1309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
talnufai wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(write_io): capture complete generate() I/O traces for LLM and VLM models #1309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
52b1259
fix: capture correct chunked inputs and all iterations in write_io_files
talnufai Aug 20, 2026
6959ec2
fix(write_io_files): capture vision encoder I/O and fix subdir naming…
talnufai Aug 20, 2026
d0b1b0e
fix(write_io): timestamp io_dir per generate() call and fix dual-QPC …
talnufai Aug 20, 2026
3a6f30a
Lint & Format
talnufai Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions QEfficient/generation/text_generation_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -848,7 +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(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,
Expand Down Expand Up @@ -1020,8 +1020,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)
Expand Down Expand Up @@ -1055,8 +1056,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)
Expand Down
12 changes: 11 additions & 1 deletion QEfficient/generation/vlm_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -673,6 +673,16 @@ 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:
Expand Down
46 changes: 30 additions & 16 deletions QEfficient/transformers/models/modeling_auto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
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
Expand Down Expand Up @@ -110,6 +111,16 @@
}


def get_io_dir(onnx_path: str) -> str:
"""Return a timestamped io_dir path under the model's onnx directory.

Format: <onnx_dir>/io_dir/<YYYYMMDD_HHMMSS_EST>
"""
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:
"""
Resolve torch_dtype in kwargs before calling from_pretrained.
Expand Down Expand Up @@ -731,7 +742,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:
Expand Down Expand Up @@ -2292,7 +2303,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:
Expand Down Expand Up @@ -2462,6 +2473,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}
Expand Down Expand Up @@ -2565,7 +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(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
Expand Down Expand Up @@ -2624,14 +2639,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)
Expand Down Expand Up @@ -3055,7 +3071,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
Expand Down Expand Up @@ -3181,7 +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, "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"]

Expand Down Expand Up @@ -3226,8 +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, "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)
Expand Down Expand Up @@ -4806,7 +4821,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):
Expand Down Expand Up @@ -5163,7 +5178,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:
Expand Down Expand Up @@ -5195,7 +5210,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)
Expand All @@ -5213,8 +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, "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)
Expand Down Expand Up @@ -5464,7 +5478,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:
Expand Down