From cadf74846d18e1e1727f6c87102711428b7a8b85 Mon Sep 17 00:00:00 2001 From: Kushal Dulla Date: Tue, 1 Sep 2026 22:11:43 +0530 Subject: [PATCH 1/3] allowed prefill blocking modes for all models Signed-off-by: Kushal Dulla --- QEfficient/blocking/attention_blocking.py | 5 ++--- .../models/qwen3_vl_moe/modeling_qwen3_vl_moe.py | 1 - 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/QEfficient/blocking/attention_blocking.py b/QEfficient/blocking/attention_blocking.py index 8b69d870fd..88967b0907 100644 --- a/QEfficient/blocking/attention_blocking.py +++ b/QEfficient/blocking/attention_blocking.py @@ -243,9 +243,8 @@ def generic_blocked_attention_interface( prefill_only: bool = False, **kwargs, ): - strategy = _STRATEGIES[ - BlockingMode.get_final_mode(blocking_config, prefill_only=prefill_only, is_mla=is_mla, mla_kwargs=mla_kwargs) - ] + prefill_only = prefill_only or blocking_config.mode.is_prefill + strategy = _STRATEGIES[BlockingMode.get_final_mode(blocking_config, is_mla=is_mla, mla_kwargs=mla_kwargs)] cache_kwargs = {"position_ids": position_ids, "batch_index": batch_index} diff --git a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index c36554ca38..bb2c210e41 100644 --- a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -452,7 +452,6 @@ def forward( batch_index=batch_index, position_ids=position_ids[0], past_seen_tokens=past_seen_tokens, - prefill_only=blocking_config.mode.is_prefill, ) else: key_states, value_states, attention_mask, _ = past_key_value_update( From 66b459b1a6e501f53acbea006f7fc4e9d35eff1c Mon Sep 17 00:00:00 2001 From: Kushal Dulla Date: Tue, 1 Sep 2026 22:38:23 +0530 Subject: [PATCH 2/3] formatting fix Signed-off-by: Kushal Dulla --- QEfficient/blocking/attention_blocking.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/QEfficient/blocking/attention_blocking.py b/QEfficient/blocking/attention_blocking.py index 88967b0907..d40cc727d9 100644 --- a/QEfficient/blocking/attention_blocking.py +++ b/QEfficient/blocking/attention_blocking.py @@ -244,7 +244,9 @@ def generic_blocked_attention_interface( **kwargs, ): prefill_only = prefill_only or blocking_config.mode.is_prefill - strategy = _STRATEGIES[BlockingMode.get_final_mode(blocking_config, is_mla=is_mla, mla_kwargs=mla_kwargs)] + strategy = _STRATEGIES[ + BlockingMode.get_final_mode(blocking_config, prefill_only=prefill_only, is_mla=is_mla, mla_kwargs=mla_kwargs) + ] cache_kwargs = {"position_ids": position_ids, "batch_index": batch_index} From 65e562bc4662d7e2e4ce4160a46194ca42f4832e Mon Sep 17 00:00:00 2001 From: Kushal Dulla Date: Thu, 10 Sep 2026 00:37:33 +0530 Subject: [PATCH 3/3] Qwen3.5 cache changes for batch fold blocking Signed-off-by: Kushal Dulla --- .../transformers/models/modeling_auto.py | 5 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 79 ++- .../qwen3_5_moe/qwen3_5_moe_blocking.py | 491 ++++++++++++++++++ 3 files changed, 558 insertions(+), 17 deletions(-) create mode 100644 examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_moe_blocking.py diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index cea08f6663..03e5cdb157 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1608,7 +1608,10 @@ def export( "continuous_batching": self.continuous_batching, "comp_ctx_lengths": self.comp_ctx_lengths_decode, } - if getattr(self.model.config, "model_type", None) == "qwen3_vl_moe": + if ( + getattr(self.model.config, "model_type", None) == "qwen3_vl_moe" + or getattr(self.model.config, "model_type", None) == "qwen3_5_moe" + ): _blocking_cfg = self.lang_model.hash_params.get("blocking_kwargs", None) batch_fold = ( not prefill_only and _blocking_cfg is not None and _blocking_cfg.mode == BlockingMode.KV_BATCH_FOLD diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index d6a04adb4a..b3dcf9da4e 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -211,12 +211,55 @@ def read_only_blockedKV(self, start_index: int, end_index: int, layer_idx: int, raise ValueError(f"Layer {layer_idx} is not a full_attention layer") return layer.read_only_blockedKV(start_index, end_index, cache_kwargs) + def read_only_blocked_K(self, start_index: int, end_index: int, layer_idx: int, cache_kwargs: dict): + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.read_only_blocked_K(start_index, end_index, cache_kwargs) + + def read_only_blocked_V(self, start_index: int, end_index: int, layer_idx: int, cache_kwargs: dict): + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.read_only_blocked_V(start_index, end_index, cache_kwargs) + + def read_only_blocked_K_batch( + self, start_index: int, end_index: int, layer_idx: int, cache_kwargs: dict, folded_cache=None + ): + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.read_only_blocked_K_batch(start_index, end_index, cache_kwargs, folded_cache=folded_cache) + + def read_only_blocked_V_batch( + self, start_index: int, end_index: int, layer_idx: int, cache_kwargs: dict, folded_cache=None + ): + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.read_only_blocked_V_batch(start_index, end_index, cache_kwargs, folded_cache=folded_cache) + def write_only(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: dict): layer = self.kv_layers[layer_idx] if layer is None: raise ValueError(f"Layer {layer_idx} is not a full_attention layer") return layer.write_only(key_states, value_states, cache_kwargs) + def get_batch_folded_kv(self, layer_idx): + """Return retained K/V cache for one layer in batch-folded compute layout.""" + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.get_batch_folded_kv() + + def write_only_batch( + self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: dict + ): + layer = self.kv_layers[layer_idx] + if layer is None: + raise ValueError(f"Layer {layer_idx} is not a full_attention layer") + return layer.write_only_batch(key_states, value_states, cache_kwargs) + def has_previous_state(self, layer_idx=None) -> bool: if layer_idx is not None: if layer_idx < 0 or layer_idx >= len(self.layer_types): @@ -2033,10 +2076,15 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): return lang, compiler_options def get_onnx_dynamic_axes( - self, comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False + self, + comp_ctx_lengths: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + batch_fold: bool = False, ): num_layers = self.config.text_config.num_hidden_layers batch_axis_name = "full_batch_size" if continuous_batching else "batch_size" + input_batch_axis_name = "full_batch_size" if continuous_batching and batch_fold else "batch_size" vision_dynamic_axes = { "pixel_values": {0: "grid_height", 1: "grid_width"}, @@ -2044,8 +2092,8 @@ def get_onnx_dynamic_axes( } lang_dynamic_axes = { - "input_ids": {0: "batch_size", 1: "seq_len"}, - "position_ids": {1: "batch_size", 2: "seq_len"}, + "input_ids": {0: input_batch_axis_name, 1: "seq_len"}, + "position_ids": {1: input_batch_axis_name, 2: "seq_len"}, "vision_embeds": {0: "vision_batch_size", 1: "vision_size"}, } @@ -2058,7 +2106,7 @@ def get_onnx_dynamic_axes( lang_dynamic_axes[f"recurrent_state.{i}"] = {0: batch_axis_name} if continuous_batching: - lang_dynamic_axes["batch_index"] = {0: "batch_size"} + lang_dynamic_axes["batch_index"] = {0: input_batch_axis_name} if comp_ctx_lengths is not None: lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} @@ -2081,25 +2129,31 @@ def get_dummy_inputs( continuous_batching: bool = False, **kwargs, ): + batch_fold = kwargs.pop("batch_fold", False) + bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) + fbs = constants.ONNX_EXPORT_EXAMPLE_FBS + if continuous_batching and batch_fold: + bs = fbs + inputs_shapes = {} dummy_seq_len = 32 - inputs_shapes["input_ids"] = (constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, dummy_seq_len) + inputs_shapes["input_ids"] = (bs, dummy_seq_len) inputs_shapes["position_ids"] = ( 4, - constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, + bs, dummy_seq_len, ) inputs_shapes["pixel_values"] = (11008, 1536) inputs_shapes["image_grid_thw"] = ( - constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, + bs, 1, 86, 128, ) inputs_shapes["vision_embeds"] = ( - constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, + bs, 2752, self.model.config.text_config.hidden_size, ) @@ -2114,19 +2168,12 @@ def get_dummy_inputs( lang_inputs["input_ids"] = torch.zeros((inputs_shapes["input_ids"]), dtype=torch.int64) lang_inputs["vision_embeds"] = torch.zeros((inputs_shapes["vision_embeds"]), dtype=float_dtype) lang_inputs["position_ids"] = ( - ( - torch.arange(dummy_seq_len, dtype=torch.int64) - .view(1, dummy_seq_len) - .repeat(constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, 1) - ) + (torch.arange(dummy_seq_len, dtype=torch.int64).view(1, dummy_seq_len).repeat(bs, 1)) .unsqueeze(0) .repeat(4, 1, 1) ) lang_inputs["image_idx"] = torch.zeros((inputs_shapes["image_idx"]), dtype=torch.int64) - bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE - fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS - kv_cache_shape = get_padding_shape_from_config( config=self.model.config.text_config, batch_size=fbs if continuous_batching else bs, diff --git a/examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_moe_blocking.py b/examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_moe_blocking.py new file mode 100644 index 0000000000..9755fff61e --- /dev/null +++ b/examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_moe_blocking.py @@ -0,0 +1,491 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import functools +import os +from time import perf_counter + +import numpy as np +import requests +import torch +import transformers +from PIL import Image +from qwen_vl_utils import process_vision_info +from transformers import AutoConfig, AutoProcessor + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession + + +def _parse_args(): + parser = argparse.ArgumentParser( + description="Compile or reuse Qwen3.5 MoE vision, prefill, and decode QPCs, then run generation." + ) + parser.add_argument("--model-id", default="Qwen/Qwen3.6-35B-A3B") + parser.add_argument( + "--prefill-qpc-path", + default=os.environ.get("QEFF_PREFILL_QPC_PATH"), + help="Existing language-prefill QPC directory. If omitted, it is compiled.", + ) + parser.add_argument( + "--decode-qpc-path", + default=os.environ.get("QEFF_DECODE_QPC_PATH"), + help="Existing language-decode QPC directory. If omitted, it is compiled.", + ) + vision_group = parser.add_mutually_exclusive_group() + vision_group.add_argument( + "--vision-qpc-path", + default=os.environ.get("QEFF_VISION_QPC_PATH"), + help="Existing vision QPC directory. If omitted, it is compiled unless --skip-vision is set.", + ) + vision_group.add_argument( + "--skip-vision", + action="store_true", + help="Run a text-only prompt and neither compile nor load a vision QPC.", + ) + parser.add_argument("--generation-len", type=int, default=256) + parser.add_argument( + "--decode-num-devices", + type=int, + default=int(os.environ.get("QEFF_DECODE_NUM_DEVICES", "2")), + ) + return parser.parse_args() + + +def _validate_qpc_path(path, component): + path = os.path.abspath(os.path.expanduser(os.fspath(path))) + if not os.path.isfile(os.path.join(path, "programqpc.bin")): + raise FileNotFoundError(f"{component} QPC does not contain programqpc.bin: {path}") + return path + + +def _compiled_qpc_path(compile_result, key, component): + path = compile_result.get(key) if isinstance(compile_result, dict) else compile_result + if not path: + raise RuntimeError(f"Compilation did not return {key!r}: {compile_result!r}") + return _validate_qpc_path(path, component) + + +args = _parse_args() +model_id = args.model_id +DECODE_NUM_DEVICES = args.decode_num_devices +config = AutoConfig.from_pretrained(model_id) + +# For faster execution user can run with lesser layers, For Testing Purpose Only +config.vision_config.depth = 5 +config.text_config.num_hidden_layers = 1 +config.torch_dtype = "float16" +layer_types = list(getattr(config.text_config, "layer_types", [])) +if len(layer_types) < config.text_config.num_hidden_layers: + layer_types.extend(["full_attention"] * (config.text_config.num_hidden_layers - len(layer_types))) +config.text_config.layer_types = layer_types[: config.text_config.num_hidden_layers] + + +def _resolve_retained_state(source_outputs, logical_name): + suffixes = ("_InternalRetainedState", "_RetainedState") + matches = [] + for output_name, value in source_outputs.items(): + if not (output_name == f"{logical_name}_RetainedState" or output_name.startswith(f"{logical_name}_")): + continue + for suffix in suffixes: + if output_name.endswith(suffix): + matches.append((output_name[: -len(suffix)], value)) + break + if len(matches) != 1: + available = sorted(name for name in source_outputs if name.startswith(logical_name)) + raise KeyError(f"Expected one retained-state output for {logical_name!r}, found {available or 'none'}") + return matches[0] + + +def _update_retained_states(target_inputs, source_outputs): + for layer_idx, layer_type in enumerate(config.text_config.layer_types): + state_names = ( + (f"past_key.{layer_idx}", f"past_value.{layer_idx}") + if layer_type == "full_attention" + else (f"conv_state.{layer_idx}", f"recurrent_state.{layer_idx}") + ) + for logical_name in state_names: + target_name, value = _resolve_retained_state(source_outputs, logical_name) + target_inputs[target_name] = value + + +PREFILL_SEQ_LEN = 1024 +GDN_CHUNK_SIZE = 1024 +CTX_LEN = 14 * 1024 +BATCH_SIZE = 8 # Per-slot prefill batch size +BS = BATCH_SIZE +FULL_BATCH_SIZE = BATCH_SIZE # Total concurrent CB slots + + +def _resize_gdn_chunk_buffers(linear_attn, chunk_size: int): + device = linear_attn._mask_causal.device + linear_attn.register_buffer( + "_mask_causal", + torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=0), + persistent=False, + ) + linear_attn.register_buffer( + "_mask_strict", + torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=device), diagonal=1), + persistent=False, + ) + linear_attn.register_buffer( + "_ones_lower", + torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.float32, device=device), diagonal=0), + persistent=False, + ) + linear_attn.register_buffer( + "_eye", + torch.eye(chunk_size, dtype=torch.float32, device=device), + persistent=False, + ) + + +def apply_chunk_size(qeff_model, chunk_size: int): + """Bind each GDN layer to script-local chunk masks sized for this export.""" + language_model = getattr(qeff_model.model, "language_model", None) + if language_model is None: + language_model = qeff_model.model.model.language_model + + for decoder_layer in language_model.layers: + linear_attn = getattr(decoder_layer, "linear_attn", None) + if linear_attn is not None: + _resize_gdn_chunk_buffers(linear_attn, chunk_size) + linear_attn.chunk_gated_delta_rule = functools.partial( + linear_attn.torch_chunk_gated_delta_rule_qeff, chunk_size=chunk_size + ) + + +qeff_model = QEFFAutoModelForImageTextToText.from_pretrained( + model_id, + attn_implementation="eager", + kv_offload=True, + config=config, +) +tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) +processor = AutoProcessor.from_pretrained(model_id) + +# Must run before qeff_model.compile() so the traced chunk size matches the +# registered _mask_causal/_mask_strict/_ones_lower/_eye buffers. +apply_chunk_size(qeff_model, GDN_CHUNK_SIZE) + +# Online-prefill KV blocking (matches GQA_Benchmark_Merged/qwen3_35b_a3b_online_prefill.sh: +# NUM_KV_BLOCKS=8, IMPLS=online, Q_BLOCK_SIZE=256 -> num_q_blocks=ceil(1024/256)=4, +# Q_HEAD_BLOCK_CHUNK=1 -> n_rep_chunk=1). +# To disable KV blocking, comment out the qaic_config= line in the prefill compile() call below. +qaic_config = { + "blocking_mode": "prefill_online", + "num_kv_blocks": 16, + "num_q_blocks": 4, + "n_rep_chunk": 1, + "skip_kv": True, +} + +# Decode-time KV blocking: + +decode_qaic_config = { + "blocking_mode": "kv_batch_fold", + "num_kv_blocks": 4, + "skip_kv": True, +} + +enable_blocking = True + +generation_len = args.generation_len +# os.environ["QAIC_COMPILER_OPTS_UNSUPPORTED"] = ( +# "-aic-user-order -aic-hoist-vtcm-loads=false " +# "-aic-op-stats-verbosity 2 -aic-hmx-async=0 -aic-userdma-async=0 -aic-sync-ts-starts" +# ) +os.environ["QAIC_COMPILER_OPTS_UNSUPPORTED"] = "-aic-op-stats-verbosity 2 -aic-hoist-vtcm-loads=false " +skip_vision = args.skip_vision + +if skip_vision: + vision_qpc_path = None +elif args.vision_qpc_path: + vision_qpc_path = _validate_qpc_path(args.vision_qpc_path, "vision") + print(f"Using compiled vision QPC: {vision_qpc_path}") +else: + vision_compile_result = qeff_model.compile( + batch_size=BS, + full_batch_size=FULL_BATCH_SIZE, + prefill_seq_len=PREFILL_SEQ_LEN, + ctx_len=CTX_LEN, + height=354, + width=536, + num_cores=16, + num_devices=1, + mos=1, + mxfp6_matmul=True, + aic_enable_depth_first=True, + skip_vision=skip_vision, + split_model_io=True, + skip_lang=True, + use_onnx_subfunctions=True, + ) + vision_qpc_path = _compiled_qpc_path(vision_compile_result, "vision_qpc_path", "vision") + print(f"Compiled vision QPC: {vision_qpc_path}") + +skip_prefill = True +if not skip_prefill: + if args.prefill_qpc_path: + prefill_qpc_path = _validate_qpc_path(args.prefill_qpc_path, "prefill") + print(f"Using compiled prefill QPC: {prefill_qpc_path}") + else: + prefill_compile_result = qeff_model.compile( + batch_size=BS, + kv_cache_batch_size=FULL_BATCH_SIZE, + full_batch_size=BS, + prefill_seq_len=PREFILL_SEQ_LEN, + ctx_len=CTX_LEN, + height=354, + width=536, + num_cores=16, + num_devices=1, + mxfp6_matmul=True, + mxint8_kv_cache=False, + retain_full_kv=True, + split_model_io=True, # This should be used for disagg serving via VLLM + # mos=1, + user_tiled=True, + aic_enable_depth_first=False, + prefill_only=True, + enable_chunking=True, + skip_vision=True, + use_onnx_subfunctions=True, + stats_level=50, + ddr_stats=True, + aic_pmu_recipe="KernelUtil", + aic_perf_metrics=True, + qaic_config=qaic_config, # Enable KV blocking - comment out to disable + kv_cache_prefix="vllmKvCache", + # custom_IO_list_file="/local/mnt/workspace/mkshirsa/qwen_36/efficient-transformers/Qwen3_5MoeForConditionalGeneration/Qwen3_5MoeDecoderWrapper-3f73445da9adc289/qpc-37d931bb34ea0c73/custom_io_kvint8.yaml" + ) + prefill_qpc_path = _compiled_qpc_path(prefill_compile_result, "lang_prefill_qpc_path", "prefill") + print(f"Compiled prefill QPC: {prefill_qpc_path}") + +if args.decode_qpc_path: + decode_qpc_path = _validate_qpc_path(args.decode_qpc_path, "decode") + print(f"Using compiled decode QPC: {decode_qpc_path}") +else: + print("Decode compiling") + decode_compile_result = qeff_model.compile( + batch_size=BS, + # kv_cache_batch_size=FULL_BATCH_SIZE, + full_batch_size=FULL_BATCH_SIZE, + ctx_len=CTX_LEN, + height=354, + width=536, + num_cores=16, + num_devices=DECODE_NUM_DEVICES, + mxfp6_matmul=True, + mxint8_kv_cache=True, + retain_full_kv=True, + split_model_io=True, # This should be used for disagg serving via VLLM + aic_enable_depth_first=False, + user_tiled=True, + prefill_only=False, + skip_vision=True, + use_onnx_subfunctions=True, + stats_level=70, + ddr_stats=True, + aic_pmu_recipe="KernelUtil", + aic_perf_metrics=True, + qaic_config=decode_qaic_config, # Enable KV blocking - comment out to disable + kv_cache_prefix="vllmKvCache", + # network_specialization_config="/path/to/specializations.json" + ) + decode_qpc_path = _compiled_qpc_path(decode_compile_result, "lang_decode_qpc_path", "decode") + print(f"Compiled decode QPC: {decode_qpc_path}") + +if enable_blocking: + print("\n" + "=" * 80) + print("Verifying KV Blocking Applied During Compilation") + print("=" * 80) + + # The compile() method internally calls BlockingAttentionTransform.apply() + # which sets attn_blocking_config on all supported attention modules + # This happens BEFORE ONNX export, so blocking operations are in the ONNX graph + + if qaic_config and qaic_config.get("blocking_mode"): + print("āœ“ qaic_config passed to compile():") + print(f" Blocking Mode: {qaic_config.get('blocking_mode')}") + print(f" Num KV Blocks: {qaic_config.get('num_kv_blocks')}") + print(f" Skip KV: {qaic_config.get('skip_kv', False)}") + print("\nāœ“ BlockingAttentionTransform.apply() called during compile()") + print(" - Sets attn_blocking_config on all supported attention modules") + print(" - Blocked attention forward pass is used during ONNX export") + print(" - Blocking operations are in the ONNX graph and QPC") + print("\n Status: ACTIVE") + print(" Verification: Config-based verification") + print(" Note: Blocking IS applied - torch model is freed after ONNX export") + else: + print("āœ— No qaic_config provided - eager attention will be used") + print(" Status: INACTIVE - Model compiled without blocking") + + print("=" * 80 + "\n") + +lang_prefill_session = QAICInferenceSession(prefill_qpc_path) +lang_decode_session = QAICInferenceSession(decode_qpc_path) + +if skip_vision: + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Tell me about yourself."}, + ], + }, + ] +else: + ### IMAGE + TEXT ### + image_url = "https://picsum.photos/id/237/536/354" + image = Image.open(requests.get(image_url, stream=True).raw) + + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "Describe all the colors seen in the image."}, + ], + }, + ] + vision_session = QAICInferenceSession(vision_qpc_path) + + +messages = [messages] * BS + +texts = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in messages] + +image_inputs, video_inputs = process_vision_info(messages) +inputs = processor( + text=texts, + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", +) + + +inputs = qeff_model.model.prepare_inputs_for_generation(inputs=inputs, prefill_seq_len=PREFILL_SEQ_LEN, batch_size=BS) + +pad_token_id = 1 +input_len = inputs["attention_mask"].sum(1, keepdims=True) +input_ids_length = inputs["input_ids"].shape[1] +num_chunks = -(input_ids_length // -PREFILL_SEQ_LEN) # ceil divide without float +padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len + +print(f"generation_len : {generation_len}") +generated_ids = np.full((BS, generation_len + 1), pad_token_id) + + +inputs["input_ids"] = torch.nn.functional.pad( + inputs["input_ids"], + (0, padded_len - input_ids_length), + "constant", + pad_token_id, +) +inputs["attention_mask"] = torch.nn.functional.pad( + inputs["attention_mask"], (0, padded_len - input_ids_length), "constant", 0 +) + +for k, v in inputs.items(): + inputs[k] = np.array(v) + + +vision_inputs = { + k: v + for k, v in inputs.items() + if k in {"pixel_values", "image_masks", "image_input_idx", "valid_idx", "aspect_ratio_ids", "aspect_ratio_mask"} +} + +vision_inputs_fp16 = {"pixel_values", "image_masks"} +vision_inputs.update({k: vision_inputs[k].astype("float16") for k in vision_inputs_fp16 if k in vision_inputs}) + +vision_start = perf_counter() +vision_outputs = {} +if vision_inputs: + vision_outputs = vision_session.run(vision_inputs) +vision_end = perf_counter() + +lang_inputs = {k: v for k, v in inputs.items() if k not in vision_inputs} +if "position_ids" in inputs: + lang_inputs["position_ids"] = inputs["position_ids"] + lang_inputs.pop("attention_mask") +else: + lang_inputs["position_ids"] = np.where( + lang_inputs.pop("attention_mask"), np.arange(padded_len), -1 + ) # Need to use -1 as position_ids for invalid tokens + +lang_inputs["image_idx"] = np.array([[0]]) + +if not skip_vision: + lang_inputs["vision_embeds"] = vision_outputs["vision_embeds"] + +# RUN prefill +lang_start = perf_counter() +lang_prefill_session.set_buffers(vision_outputs) + +all_outputs = [] +chunk_inputs = lang_inputs.copy() +for i in range(num_chunks): + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][..., i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] + outputs = lang_prefill_session.run(chunk_inputs) + _update_retained_states(chunk_inputs, outputs) + chunk_inputs["image_idx"] = outputs["image_idx_output"] +prefill_time = perf_counter() - lang_start + vision_end - vision_start +print(f"Prefill time : {prefill_time:.2f} secs") + +all_outputs.append(np.argmax(outputs["logits"])) +decode_inputs = { + "input_ids": np.argmax(outputs["logits"]).reshape(1, 1), + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True) + 1, +} + +_update_retained_states(decode_inputs, outputs) + +decode_inputs["image_idx"] = outputs["image_idx_output"] + +if not skip_vision: + decode_inputs["vision_embeds"] = outputs["vision_embeds_RetainedState"] + +st = perf_counter() +decode_out = lang_decode_session.run(decode_inputs) +print(f"time for first run of decode with KV as input = {perf_counter() - st} sec\n") + +all_outputs.append(np.argmax(decode_out["logits"])) +pos_id = np.max(decode_inputs["position_ids"], axis=-1, keepdims=True) + 1 +loop_decode_inputs = { + "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), + "position_ids": pos_id, +} + +_update_retained_states(loop_decode_inputs, decode_out) + +loop_decode_inputs["image_idx"] = decode_out["image_idx_output"] + +if not skip_vision: + loop_decode_inputs["vision_embeds"] = decode_out["vision_embeds_RetainedState"] + + +st = perf_counter() +for i in range(generation_len - 2): + decode_out = lang_decode_session.run(loop_decode_inputs) + all_outputs.append(np.argmax(decode_out["logits"])) + pos_id += 1 + _update_retained_states(loop_decode_inputs, decode_out) + loop_decode_inputs.update( + { + "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), + "position_ids": pos_id, + } + ) +ft = perf_counter() +print(f"decode tok/sec={(generation_len - 2) / (ft - st)}") +print(f"\noutput\n{tokenizer.decode(all_outputs)}")