From 51c6f402446cf8b52e832e6f0826953ec1b26e15 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Wed, 2 Sep 2026 15:23:52 +0530 Subject: [PATCH 1/3] move pytorch transforms Signed-off-by: Mamta Singh --- QEfficient/base/modeling_qeff.py | 259 ++++++++++++++- .../exporter/export_hf_to_cloud_ai_100.py | 3 + .../transformers/models/modeling_auto.py | 301 ++++++------------ QEfficient/utils/_utils.py | 2 +- .../test_automodel_for_causal_lm.py | 11 +- .../test_automodel_for_speech_seq2seq.py | 11 +- 6 files changed, 361 insertions(+), 226 deletions(-) diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 77209b94b1..ca5181060d 100644 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -40,8 +40,24 @@ from QEfficient.generation.cloud_infer import QAICInferenceSession from QEfficient.transformers.models.pytorch_transforms import ( BlockingAttentionTransform, + CustomOpsTransform, + KVCacheExternalModuleMapperTransform, + KVCacheTransform, OptimizedMoETransform, + PoolingTransform, + PrefillOnlyChunkedTransform, + PrefillOnlyExternalModuleMapperTransform, + PrefillOnlyTransform, ReplicateKVHeadTransform, + RevertPrefillKeepAttentionTransform, + RevertPrefillOnlyExternalModuleMapperTransform, + RevertPrefillOnlyTransform, + SamplerTransform, + SimpleDecodeMoeTransform, + SpDTransform, + TextClassificationTransform, + VlmKVOffloadTransform, + VlmNoKVOffloadTransform, ) from QEfficient.utils import ( align_kv_input_names_to_retained_outputs, @@ -67,6 +83,47 @@ "moe_prefill_packed_chunk_size is no longer supported; use qaic_config['moe_config']['expert_parallel_chunk_size']" ) +_SPECIALIZED_DISAGG_SERVING_MODEL_ARCH = {"gpt_oss", "qwen3_moe", "glm4_moe", "kimi_k2", "kimi_k25", "gemma4"} + +_MODELING_AUTO_PYTORCH_TRANSFORMS = { + "QEFFAutoModel": ((CustomOpsTransform,), ()), + "QEFFAutoModelForSequenceClassification": ((CustomOpsTransform, TextClassificationTransform), ()), + "QEffVisionEncoderForTextImageToTextModel": ( + (), + (CustomOpsTransform, KVCacheTransform, KVCacheExternalModuleMapperTransform), + ), + "QEffCausalLMForTextImageToTextModel": ( + (), + (CustomOpsTransform, KVCacheTransform, VlmKVOffloadTransform, SimpleDecodeMoeTransform), + ), + "_QEFFAutoModelForImageTextToTextSingleQPC": ( + (), + ( + CustomOpsTransform, + KVCacheTransform, + KVCacheExternalModuleMapperTransform, + VlmNoKVOffloadTransform, + SimpleDecodeMoeTransform, + ), + ), + "QEFFAutoModelForCausalLM": ( + (), + (CustomOpsTransform, KVCacheTransform, KVCacheExternalModuleMapperTransform, SimpleDecodeMoeTransform), + ), + "QEFFAutoModelForSpeechSeq2Seq": ((CustomOpsTransform,), (KVCacheTransform,)), + "QEFFAutoModelForCTC": ((CustomOpsTransform,), ()), +} + +_PREFILL_TRANSFORM_MODEL_CLASSES = { + "QEFFAutoModelForCausalLM", + "QEffCausalLMForTextImageToTextModel", + "_QEFFAutoModelForImageTextToTextSingleQPC", +} + +_SPD_TRANSFORM_MODEL_CLASSES = {"QEFFAutoModelForCausalLM"} + +_SAMPLER_TRANSFORM_MODEL_CLASSES = {"QEFFAutoModelForCausalLM", "QEffCausalLMForTextImageToTextModel"} + def reject_legacy_moe_prefill_packed_chunk_size(kwargs: Optional[dict]) -> None: if kwargs and "moe_prefill_packed_chunk_size" in kwargs: @@ -118,7 +175,7 @@ class QEFFBaseModel(ABC): Provides certain utility methods to be used by child classes. Class variables: - :_pytorch_transforms: Pytorch transformations to be applied after initialization. + :_pytorch_transforms: Pytorch transformations to be applied before export/compile. :_onnx_transforms: ONNX transformations to be applied after ONNX export. """ @@ -126,11 +183,152 @@ class QEFFBaseModel(ABC): _end = 0 _total_layers = None _layerwise_active = False - _pytorch_transforms: List[PytorchTransform] + _pytorch_transforms: List[PytorchTransform] = [] _onnx_transforms = [BaseOnnxTransform] def _transform_names(self) -> List[str]: - return [x.__name__ for x in self._pytorch_transforms + self._onnx_transforms] + return [x.__name__ for x in self._all_pytorch_transforms() + self._onnx_transforms] + + def _modeling_auto_pytorch_transform_groups(self): + for cls in type(self).mro(): + if transforms := _MODELING_AUTO_PYTORCH_TRANSFORMS.get(cls.__name__): + pre_quant_transforms, post_quant_transforms = transforms + return list(pre_quant_transforms), list(post_quant_transforms) + return [], [] + + def _all_pytorch_transforms(self) -> List[PytorchTransform]: + pre_quant_transforms, post_quant_transforms = self._modeling_auto_pytorch_transform_groups() + class_transforms = [] + proxy_transforms = [] + for transform in self._pytorch_transforms: + if transform.__name__ == "QeffProxyModuleTransform": + proxy_transforms.append(transform) + else: + class_transforms.append(transform) + transforms = [] + for transform in pre_quant_transforms + class_transforms + post_quant_transforms + proxy_transforms: + if transform not in transforms: + transforms.append(transform) + return transforms + + def _apply_pytorch_transforms(self) -> bool: + any_transformed = False + for transform in self._all_pytorch_transforms(): + self.model, transformed = transform.apply(self.model) + any_transformed = any_transformed or transformed + return any_transformed + + def _post_pytorch_transform(self) -> bool: + return False + + def _apply_pooling_transform(self, pooling=None) -> None: + if pooling: + self.model, _ = PoolingTransform.apply(self.model, pooling) + + def _apply_prefill_transform( + self, + enable: Optional[bool] = True, + enable_chunking: Optional[bool] = False, + retain_full_kv: Optional[bool] = False, + use_external_module_mapper: Optional[bool] = False, + ) -> None: + if enable: + if use_external_module_mapper: + self.model, _ = PrefillOnlyExternalModuleMapperTransform.apply(self.model) + if enable_chunking: + self.model, _ = PrefillOnlyChunkedTransform.apply(self.model) + else: + self.model, _ = PrefillOnlyTransform.apply(self.model) + else: + if use_external_module_mapper: + self.model, _ = RevertPrefillOnlyExternalModuleMapperTransform.apply(self.model) + if retain_full_kv: + self.model, _ = RevertPrefillKeepAttentionTransform.apply(self.model) + else: + self.model, _ = RevertPrefillOnlyTransform.apply(self.model) + + def _uses_external_prefill_mapper(self) -> bool: + return KVCacheExternalModuleMapperTransform in self._all_pytorch_transforms() + + def _supports_prefill_transform_options(self) -> bool: + return any(cls.__name__ in _PREFILL_TRANSFORM_MODEL_CLASSES for cls in type(self).mro()) + + def _should_apply_prefill_transform_from_options(self) -> bool: + class_names = {cls.__name__ for cls in type(self).mro()} + if "_QEFFAutoModelForImageTextToTextSingleQPC" in class_names: + return True + model_type = getattr(getattr(self, "config", None), "model_type", None) or getattr( + getattr(self.model, "config", None), "model_type", None + ) + return model_type in _SPECIALIZED_DISAGG_SERVING_MODEL_ARCH + + def _supports_spd_transform(self) -> bool: + return any(cls.__name__ in _SPD_TRANSFORM_MODEL_CLASSES for cls in type(self).mro()) + + def _supports_sampler_transform(self) -> bool: + return any(cls.__name__ in _SAMPLER_TRANSFORM_MODEL_CLASSES for cls in type(self).mro()) + + def _update_prefill_hash_params( + self, + *, + prefill_only: Optional[bool], + enable_chunking: Optional[bool], + retain_full_kv: Optional[bool], + ) -> None: + model_type = getattr(getattr(self.model, "config", None), "model_type", None) + architectures = getattr(getattr(self.model, "config", None), "architectures", None) or [] + + if model_type in _SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: + if prefill_only and "DeepseekV3ForCausalLM" not in architectures: + self.hash_params.pop("retain_full_kv", None) + self.hash_params["prefill_only"] = True + if enable_chunking: + self.hash_params["chunking"] = True + else: + self.hash_params.pop("prefill_only", None) + self.hash_params.pop("NUM_Q_BLOCKS", None) + self.hash_params.pop("NUM_FFN_BLOCKS", None) + self.hash_params.pop("ENABLE_OPT_SWA", None) + self.hash_params.pop("chunking", None) + self.hash_params.pop("chunking_seq_len", None) + if retain_full_kv: + self.hash_params["retain_full_kv"] = True + else: + self.hash_params.pop("retain_full_kv", None) + elif prefill_only is not None: + self.hash_params["prefill_only"] = bool(prefill_only) + + def _apply_prefill_transform_from_options( + self, + *, + prefill_only: Optional[bool] = False, + prefill_seq_len: Optional[int] = None, + enable_chunking: Optional[bool] = False, + retain_full_kv: Optional[bool] = False, + ) -> None: + if prefill_only: + assert prefill_seq_len is None or prefill_seq_len > 1 + if not enable_chunking and getattr(self, "continuous_batching", False): + raise NotImplementedError( + "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" + ) + self._apply_prefill_transform( + enable=True, + enable_chunking=enable_chunking, + use_external_module_mapper=self._uses_external_prefill_mapper(), + ) + else: + self._apply_prefill_transform( + enable=False, + retain_full_kv=retain_full_kv, + use_external_module_mapper=self._uses_external_prefill_mapper(), + ) + + self._update_prefill_hash_params( + prefill_only=prefill_only, + enable_chunking=enable_chunking, + retain_full_kv=retain_full_kv, + ) def maybe_apply_replicate_kv_transform(self, model_config, num_devices: int, qaic_config: Optional[dict]) -> int: if model_config is None or qaic_config is None or "EncoderWrapper" in self.model.__class__.__name__: @@ -182,16 +380,6 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None: self.is_transformed: bool = False self._normalize_torch_dtype() - # Apply the transformations - any_transformed = False - for transform in self._pytorch_transforms: - self.model, transformed = transform.apply(self.model) - any_transformed = any_transformed or transformed - - if not any_transformed: - warnings.warn(f"No transforms applied to model: {self.model_name}. It may be an unsupported model!") - else: - logger.info(f"Pytorch transforms applied to model: {self.model_name}") def _normalize_torch_dtype(self): """ @@ -693,10 +881,12 @@ def get_onnx_path( qaic_config=qaic_config, prefill_only=prefill_only, enable_chunking=enable_chunking, + retain_full_kv=retain_full_kv, num_cores=kwargs.get("num_cores", compiler_options.get("aic_num_cores", constants.DEFAULT_AIC_NUM_CORES)), prefill_seq_len=kwargs.get("prefill_seq_len"), **compiler_options, ) + kwargs["_qeff_skip_transform"] = True with export_from_compile(): self.export(**kwargs) @@ -945,6 +1135,41 @@ def transform( qaic_config: Optional[dict] = None, **compiler_options, ): + if not self.is_transformed: + any_transformed = self._apply_pytorch_transforms() + pooling = compiler_options.pop("pooling", getattr(self, "_pooling", None)) + if pooling: + self._apply_pooling_transform(pooling) + any_transformed = True + + any_transformed = self._post_pytorch_transform() or any_transformed + + qaic_config_for_transforms = qaic_config or getattr(self.model, "qaic_config", None) + if self._supports_spd_transform(): + self.model, spd_transformed = SpDTransform.apply( + self.model, + qaic_config_for_transforms, + **compiler_options, + ) + self.is_tlm = getattr(self, "is_tlm", False) or spd_transformed + any_transformed = any_transformed or spd_transformed + + if self._supports_sampler_transform(): + self.model, sampler_transformed = SamplerTransform.apply( + self.model, + qaic_config_for_transforms, + **compiler_options, + ) + any_transformed = any_transformed or sampler_transformed + if getattr(self, "is_tlm", False) and getattr(self.model, "qaic_config", None) is not None: + self.model.qaic_config["return_pdfs"] = True + + if not any_transformed: + warnings.warn(f"No transforms applied to model: {self.model_name}. It may be an unsupported model!") + else: + logger.info(f"Pytorch transforms applied to model: {self.model_name}") + self.is_transformed = True + # Apply the transformations that are dependent on compilation parameters model_config = getattr(self.model, "config", None) or getattr( getattr(self.model, "model", None), "config", None @@ -992,6 +1217,14 @@ def transform( hash_params=self.hash_params, ) + if self._supports_prefill_transform_options() and self._should_apply_prefill_transform_from_options(): + self._apply_prefill_transform_from_options( + prefill_only=compiler_options.get("prefill_only", False), + prefill_seq_len=compiler_options.get("prefill_seq_len", seq_len), + enable_chunking=compiler_options.get("enable_chunking", False), + retain_full_kv=compiler_options.get("retain_full_kv", False), + ) + @dump_qconfig def _compile( self, diff --git a/QEfficient/exporter/export_hf_to_cloud_ai_100.py b/QEfficient/exporter/export_hf_to_cloud_ai_100.py index 2547d9db36..620052a4ff 100644 --- a/QEfficient/exporter/export_hf_to_cloud_ai_100.py +++ b/QEfficient/exporter/export_hf_to_cloud_ai_100.py @@ -318,6 +318,9 @@ def export_lm_model_for_cloud( logger.warning(f"Overriding {onnx_dir_path}") shutil.rmtree(onnx_dir_path) + if not qeff_model.is_transformed: + qeff_model.transform(seq_len=seq_length, bs=full_batch_size or len(Constants.INPUT_STR)) + model_name = export_kvstyle_transformed_model_to_onnx( model_name=model_name, transformed_model=qeff_model.model, diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 25fac3b646..91a5796e65 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -48,24 +48,6 @@ _configure_proxy_for_model, ) from QEfficient.transformers.models.gpt_oss.modeling_gpt_oss import override_gptoss_prefill_chunking -from QEfficient.transformers.models.pytorch_transforms import ( - CustomOpsTransform, - KVCacheExternalModuleMapperTransform, - KVCacheTransform, - PoolingTransform, - PrefillOnlyChunkedTransform, - PrefillOnlyExternalModuleMapperTransform, - PrefillOnlyTransform, - RevertPrefillKeepAttentionTransform, - RevertPrefillOnlyExternalModuleMapperTransform, - RevertPrefillOnlyTransform, - SamplerTransform, - SimpleDecodeMoeTransform, - SpDTransform, - TextClassificationTransform, - VlmKVOffloadTransform, - VlmNoKVOffloadTransform, -) from QEfficient.transformers.moe.flavours import MoEFlavour from QEfficient.transformers.quantizers.auto import QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING, with_replaced_quantizers from QEfficient.transformers.quantizers.quant_transforms import ( @@ -427,7 +409,7 @@ class QEFFAutoModel(QEFFTransformersBase): """ _hf_auto_class = AutoModel - _pytorch_transforms = [CustomOpsTransform, AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] + _pytorch_transforms = [AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] # FP16Clip inlines external weights; without Split the saved protobuf exceeds 2GB for large embedders. _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform] @@ -446,12 +428,9 @@ def __init__(self, model: nn.Module, pooling=None, **kwargs): **kwargs : Additional keyword arguments passed to the base class constructor. """ + self._pooling = pooling super().__init__(model, **kwargs) - # Make Embedding specific transforms like appending pooling - if pooling: - self.model, _ = PoolingTransform.apply(self.model, pooling) - # Encoder-only models (e.g. BERT) should not be forced into cache mode. if getattr(self.model.config, "is_decoder", False) or getattr(self.model.config, "is_encoder_decoder", False): self.model.base_model.config.use_cache = True @@ -553,6 +532,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str: """ bs = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + if not kwargs.pop("_qeff_skip_transform", False): + self.transform(seq_len=seq_len, bs=bs) example_inputs = { "input_ids": torch.zeros((bs, seq_len), dtype=torch.int64), @@ -831,7 +812,7 @@ class QEFFAutoModelForSequenceClassification(QEFFTransformersBase): """ _hf_auto_class = AutoModelForSequenceClassification - _pytorch_transforms = [CustomOpsTransform, TextClassificationTransform] + _pytorch_transforms = [] _onnx_transforms = [] def __init__(self, model: nn.Module, **kwargs): @@ -924,6 +905,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str: """ bs = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + if not kwargs.pop("_qeff_skip_transform", False): + self.transform(seq_len=seq_len, bs=bs) example_inputs = { "input_ids": torch.zeros((bs, seq_len), dtype=torch.int64), @@ -1077,9 +1060,6 @@ class QEffVisionEncoderForTextImageToTextModel(QEFFBaseModel): AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, PackQuantizedInt4ToMatMulNBitsTransform, - CustomOpsTransform, - KVCacheTransform, - KVCacheExternalModuleMapperTransform, ] _onnx_transforms = [] @@ -1096,9 +1076,16 @@ def __init__(self, model: nn.modules, **kwargs): """ _configure_proxy_for_model(self, kwargs.pop("enable_proxy", False)) super().__init__(model, **kwargs) - self.model = model.get_qeff_vision_encoder() + if hasattr(model, "get_qeff_vision_encoder"): + self.model = model.get_qeff_vision_encoder() self.hash_params["qeff_auto_class"] = self.__class__.__name__ + def _post_pytorch_transform(self) -> bool: + if hasattr(self.model, "get_qeff_vision_encoder"): + self.model = self.model.get_qeff_vision_encoder() + return True + return False + def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt_weights=True, **kwargs): """ Exports the vision encoder component to ONNX format. @@ -1123,6 +1110,8 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt str Path to the generated ONNX graph file for the vision encoder. """ + if not kwargs.pop("_qeff_skip_transform", False): + self.transform() return self._export( inputs, output_names=output_names, @@ -1214,10 +1203,6 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): PackQuantizedInt4ToMatMulNBitsTransform, FP8BlockWiseDequantQwen3VLMoeTextExpertsToQwen3VLMoeTextExpertsTransform, FP8BlockWiseDequantLinearToLinearTransform, - CustomOpsTransform, - KVCacheTransform, - VlmKVOffloadTransform, - SimpleDecodeMoeTransform, ] _onnx_transforms = [] @@ -1236,33 +1221,28 @@ def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): Additional keyword arguments passed to the base class constructor. """ _configure_proxy_for_model(self, kwargs.pop("enable_proxy", False)) + self._qaic_config = qaic_config super().__init__(model, **kwargs) - self.model = model.get_qeff_language_decoder() + if hasattr(model, "get_qeff_language_decoder"): + self.model = model.get_qeff_language_decoder() self.model.qaic_config = qaic_config self.hash_params["qeff_auto_class"] = self.__class__.__name__ self.continuous_batching = False if qaic_config: if mla_absorption := qaic_config.get("mla_absorption", None): self.hash_params["mla_absorption"] = mla_absorption - setattr(self.model.language_model, "mla_absorption", mla_absorption) - - def __update_prefill_transform( - self, - enable: Optional[bool] = True, - enable_chunking: Optional[bool] = False, - retain_full_kv: Optional[bool] = False, - ): - if enable: - if enable_chunking: - self.model, tf = PrefillOnlyChunkedTransform.apply(self.model) - else: - self.model, tf = PrefillOnlyTransform.apply(self.model) - - else: - if retain_full_kv: - self.model, tf = RevertPrefillKeepAttentionTransform.apply(self.model) - else: - self.model, tf = RevertPrefillOnlyTransform.apply(self.model) + if language_model := getattr(self.model, "language_model", None): + setattr(language_model, "mla_absorption", mla_absorption) + + def _post_pytorch_transform(self) -> bool: + if hasattr(self.model, "get_qeff_language_decoder"): + self.model = self.model.get_qeff_language_decoder() + self.model.qaic_config = self._qaic_config + if self._qaic_config and (mla_absorption := self._qaic_config.get("mla_absorption", None)): + if language_model := getattr(self.model, "language_model", None): + setattr(language_model, "mla_absorption", mla_absorption) + return True + return False def export( self, @@ -1301,19 +1281,16 @@ def export( Path to the generated ONNX graph file for the language decoder. """ reject_legacy_moe_prefill_packed_chunk_size(kwargs) - if prefill_only: - assert prefill_seq_len > 1 - if not enable_chunking and self.continuous_batching: - raise NotImplementedError( - "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" - ) - self.hash_params["prefill_only"] = True - self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) - else: - self.hash_params["prefill_only"] = False - self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) - + skip_transform = kwargs.pop("_qeff_skip_transform", False) qaic_config = kwargs.pop("qaic_config", getattr(self.model, "qaic_config", None)) + if not skip_transform: + self.transform( + prefill_only=prefill_only, + enable_chunking=enable_chunking, + prefill_seq_len=prefill_seq_len, + retain_full_kv=kwargs.get("retain_full_kv", False), + qaic_config=qaic_config, + ) if QEfficient.base.modeling_qeff.QEFFBaseModel._layerwise_active: return self._export_layerwise( @@ -1457,11 +1434,6 @@ def __init__( self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.input_shapes, self.output_names = None, None - # ---Sampling--- - # Note: SamplerTransform should be applied after all other transforms - # are done. The role of the sampler is to just add nodes at the output of the - # previous transform function. - self.lang_model.model, _ = SamplerTransform.apply(self.lang_model.model, qaic_config, **kwargs) @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, qaic_config: Optional[dict] = None, **kwargs): @@ -1523,26 +1495,6 @@ def onnx_path(self): """ return [self.vision_model.onnx_path, self.lang_model.onnx_path] - def __update_prefill_transform( - self, - enable: Optional[bool] = True, - enable_chunking: Optional[bool] = False, - retain_full_kv: Optional[bool] = False, - ): - if enable: - self.model, tf = PrefillOnlyExternalModuleMapperTransform.apply(self.model) - if enable_chunking: - self.model, tf = PrefillOnlyChunkedTransform.apply(self.model) - else: - self.model, tf = PrefillOnlyTransform.apply(self.model) - - else: - self.model, tf = RevertPrefillOnlyExternalModuleMapperTransform.apply(self.model) - if retain_full_kv: - self.model, tf = RevertPrefillKeepAttentionTransform.apply(self.model) - else: - self.model, tf = RevertPrefillOnlyTransform.apply(self.model) - def export( self, export_dir: Optional[str] = None, @@ -1598,11 +1550,17 @@ def export( seq_len: int = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN qaic_config = kwargs.get("qaic_config", getattr(self.lang_model.model, "qaic_config", None)) # TODO: move this to a DA Serving utility class - if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: - if prefill_only: - self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) - else: - self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) + if not kwargs.pop("_qeff_skip_transform", False): + self.transform( + seq_len=prefill_seq_len, + num_devices=kwargs.get("num_devices", 1), + qaic_config=qaic_config, + num_cores=num_cores, + prefill_only=prefill_only, + prefill_seq_len=prefill_seq_len, + enable_chunking=enable_chunking, + retain_full_kv=kwargs.get("retain_full_kv", False), + ) onnx_kwargs = {"prefill_seq_len": seq_len, "batch_size": bs} dynamic_axes_kwargs = { "kv_offload": True, @@ -1662,6 +1620,7 @@ def export( export_dir=export_dir, offload_pt_weights=False, use_onnx_subfunctions=use_onnx_subfunctions, + _qeff_skip_transform=True, ) # TODO: remove the current pt weight offload capability once CustomLoader is in place @@ -1686,6 +1645,7 @@ def export( qaic_config=qaic_config, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + _qeff_skip_transform=True, ) return self.onnx_path @@ -2005,10 +1965,6 @@ def compile( kv_cache_batch_size = kv_cache_batch_size or full_batch_size or batch_size kv_cache_prefix = validate_kv_cache_prefix(kv_cache_prefix) - output_names = self.model.get_output_names(kv_offload=True) - # Prefix only the language-side KV-cache retained buffers (vision buffers are untouched) so the - # derived custom_io_lang keys match the prefixed names written into the exported graph. - output_names = apply_kv_cache_prefix(output_names, kv_cache_prefix) # if ccl_enabled is True read Compute-Context-Length lists if self.ccl_enabled: @@ -2033,8 +1989,15 @@ def compile( aic_num_cores=num_cores, prefill_only=prefill_only, prefill_seq_len=prefill_seq_len, + enable_chunking=enable_chunking, + retain_full_kv=compiler_options.get("retain_full_kv", False), ) + output_names = self.model.get_output_names(kv_offload=True) + # Prefix only the language-side KV-cache retained buffers (vision buffers are untouched) so the + # derived custom_io_lang keys match the prefixed names written into the exported graph. + output_names = apply_kv_cache_prefix(output_names, kv_cache_prefix) + specializations, compiler_options = self.model.get_specializations( batch_size=batch_size, prefill_seq_len=prefill_seq_len, @@ -2085,6 +2048,7 @@ def compile( _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, offload_pt_weights=offload_pt_weights, + _qeff_skip_transform=True, ) if layerwise_cache_probe: return self.lang_model.onnx_path @@ -2624,11 +2588,6 @@ class _QEFFAutoModelForImageTextToTextSingleQPC(QEFFTransformersBase, Multimodal _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, - CustomOpsTransform, - KVCacheTransform, - KVCacheExternalModuleMapperTransform, - VlmNoKVOffloadTransform, - SimpleDecodeMoeTransform, ] _onnx_transforms = [] @@ -2739,24 +2698,6 @@ def from_pretrained( **kwargs, ) - def __update_prefill_transform( - self, - enable: Optional[bool] = True, - enable_chunking: Optional[bool] = False, - retain_full_kv: Optional[bool] = False, - ): - if enable: - if enable_chunking: - self.model, tf = PrefillOnlyChunkedTransform.apply(self.model) - else: - self.model, tf = PrefillOnlyTransform.apply(self.model) - - else: - if retain_full_kv: - self.model, tf = RevertPrefillKeepAttentionTransform.apply(self.model) - else: - self.model, tf = RevertPrefillOnlyTransform.apply(self.model) - def export( self, export_dir: Optional[str] = None, @@ -2783,17 +2724,15 @@ def export( Path to the generated ONNX graph file. """ reject_legacy_moe_prefill_packed_chunk_size(kwargs) - if prefill_only: - assert prefill_seq_len > 1 - if not enable_chunking and self.continuous_batching: - raise NotImplementedError( - "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" - ) - self.hash_params["prefill_only"] = True - self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) - else: - self.hash_params["prefill_only"] = False - self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) + skip_transform = kwargs.pop("_qeff_skip_transform", False) + if not skip_transform: + self.transform( + prefill_only=prefill_only, + enable_chunking=enable_chunking, + prefill_seq_len=prefill_seq_len, + retain_full_kv=kwargs.get("retain_full_kv", False), + qaic_config=kwargs.get("qaic_config", getattr(self.model, "qaic_config", None)), + ) inputs = self.model.get_dummy_inputs(comp_ctx_lengths=self.comp_ctx_lengths_decode) dynamic_axes = self.model.get_onnx_dynamic_axes(comp_ctx_lengths=self.comp_ctx_lengths_decode) @@ -3441,10 +3380,6 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel): FP8DeQuantLinearToLinearTransform, PackQuantizedInt4ToMatMulNBitsTransform, Mxfp4GptOssExpertDequantizeTransform, - CustomOpsTransform, - KVCacheTransform, - KVCacheExternalModuleMapperTransform, - SimpleDecodeMoeTransform, ] _onnx_transforms = [] @@ -3455,39 +3390,11 @@ def prefill( enable_chunking: Optional[bool] = False, retain_full_kv: Optional[bool] = False, ): - if enable: - self.model, tf = PrefillOnlyExternalModuleMapperTransform.apply(self.model) - if enable_chunking: - self.model, tf = PrefillOnlyChunkedTransform.apply(self.model) - else: - self.model, tf = PrefillOnlyTransform.apply(self.model) - - else: - self.model, tf = RevertPrefillOnlyExternalModuleMapperTransform.apply(self.model) - if retain_full_kv: - self.model, tf = RevertPrefillKeepAttentionTransform.apply(self.model) - else: - self.model, tf = RevertPrefillOnlyTransform.apply(self.model) - - def __update_prefill_transform( - self, - enable: Optional[bool] = True, - enable_chunking: Optional[bool] = False, - retain_full_kv: Optional[bool] = False, - ): - if enable: - self.model, tf = PrefillOnlyExternalModuleMapperTransform.apply(self.model) - if enable_chunking: - self.model, tf = PrefillOnlyChunkedTransform.apply(self.model) - else: - self.model, tf = PrefillOnlyTransform.apply(self.model) - - else: - self.model, tf = RevertPrefillOnlyExternalModuleMapperTransform.apply(self.model) - if retain_full_kv: - self.model, tf = RevertPrefillKeepAttentionTransform.apply(self.model) - else: - self.model, tf = RevertPrefillOnlyTransform.apply(self.model) + self._apply_prefill_transform_from_options( + prefill_only=enable, + enable_chunking=enable_chunking, + retain_full_kv=retain_full_kv, + ) def __init__( self, @@ -3551,8 +3458,7 @@ def __init__( self.continuous_batching = continuous_batching self.model.qaic_config = qaic_config self.model.pretrained_path = kwargs.pop("pretrained_model_name_or_path", None) - self.model, transformed = SpDTransform.apply(self.model, qaic_config, **kwargs) - self.is_tlm = transformed + self.is_tlm = bool(qaic_config and qaic_config.get("speculative_model_type") is not None) self.hash_params["qeff_auto_class"] = self.__class__.__name__ self.ccl_enabled = False @@ -3564,16 +3470,6 @@ def __init__( self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.hash_params["max_seq_len_cached"] = max_seq_len_cached - # ---Sampling--- - # Note: SamplerTransform should be applied after all other transforms - # are done. The role of the sampler is to just add nodes at the output of the - # previous transform function. - self.model, transformed = SamplerTransform.apply(self.model, qaic_config, **kwargs) - # TODO : Update in qaic_config isn't updated in the hash due to SpDTransforms. Need to move - # SpDTransforms to PytorchTransforms. - if self.is_tlm: - self.model.qaic_config["return_pdfs"] = True - def __repr__(self) -> str: return self.__class__.__name__ + "\n" + self.model.__repr__() @@ -3844,31 +3740,15 @@ def export( ) ######################################## - ####### HANDLE DA PREFILL And REVERT PREFILL Transform ################ - # TODO: move this code inside self.transform in modeling_qeff.py - if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: - if prefill_only: - if not enable_chunking and self.continuous_batching: - raise NotImplementedError( - "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" - ) - self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) - self.hash_params.pop("retain_full_kv", None) - if "DeepseekV3ForCausalLM" not in (getattr(self.model.config, "architectures", None) or []): - self.hash_params["prefill_only"] = True - if enable_chunking: - self.hash_params["chunking"] = True - else: - self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) - self.hash_params.pop("prefill_only", None) - self.hash_params.pop("NUM_Q_BLOCKS", None) - self.hash_params.pop("NUM_FFN_BLOCKS", None) - self.hash_params.pop("ENABLE_OPT_SWA", None) - self.hash_params.pop("chunking", None) - self.hash_params.pop("chunking_seq_len", None) - if kwargs.get("retain_full_kv", False): - self.hash_params["retain_full_kv"] = True - ####################################################################### + if not kwargs.pop("_qeff_skip_transform", False): + self.transform( + prefill_only=prefill_only, + enable_chunking=enable_chunking, + prefill_seq_len=prefill_seq_len, + retain_full_kv=kwargs.get("retain_full_kv", False), + qaic_config=qaic_config, + num_cores=num_cores, + ) bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len: int = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN @@ -4819,7 +4699,7 @@ class QEFFAutoModelForSpeechSeq2Seq(QEFFTransformersBase, MultimodalUtilityMixin """ _hf_auto_class = AutoModelForSpeechSeq2Seq - _pytorch_transforms = [CustomOpsTransform, AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, KVCacheTransform] + _pytorch_transforms = [AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] _onnx_transforms = [] def __init__(self, model: nn.Module, **kwargs): @@ -4880,6 +4760,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str: str Path to the generated ONNX graph file. """ + if not kwargs.pop("_qeff_skip_transform", False): + self.transform() inputs = self.model.get_dummy_inputs() dynamic_axes = self.model.get_onnx_dynamic_axes() output_names = self.model.get_output_names() @@ -4971,6 +4853,7 @@ def compile( Path to the compiled QPC package. """ + self.transform(seq_len=prefill_seq_len, bs=batch_size) specializations, compiler_options = self.model.get_specializations( batch_size, encoder_ctx_len, @@ -5177,7 +5060,7 @@ class QEFFAutoModelForCTC(QEFFTransformersBase): """ _hf_auto_class = AutoModelForCTC - _pytorch_transforms = [CustomOpsTransform, AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] + _pytorch_transforms = [AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] _onnx_transforms = [] def __init__(self, model: nn.Module, **kwargs): @@ -5264,6 +5147,8 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str: """ bs = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len = constants.WAV2VEC2_MAX_SEQ_LEN + if not kwargs.pop("_qeff_skip_transform", False): + self.transform(seq_len=seq_len, bs=bs) example_inputs = { "input_values": torch.zeros((bs, seq_len), dtype=self.model.config.torch_dtype), diff --git a/QEfficient/utils/_utils.py b/QEfficient/utils/_utils.py index 5e08a68b9b..7fabdba334 100755 --- a/QEfficient/utils/_utils.py +++ b/QEfficient/utils/_utils.py @@ -802,7 +802,7 @@ def wrapper(self, *args, **kwargs): qpc_path, self.onnx_path, self.get_model_config, - [cls.__name__ for cls in self._pytorch_transforms], + [cls.__name__ for cls in self._all_pytorch_transforms()], [cls.__name__ for cls in self._onnx_transforms], kwargs.get("specializations"), kwargs.get("mdp_ts_num_devices", 1), diff --git a/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py b/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py index 9603ccc09d..bd76e56ee2 100644 --- a/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py +++ b/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py @@ -83,8 +83,9 @@ def get_auto_config_from_test_config(configs): @pytest.mark.parametrize("cb", [False, True], ids=["nocb", "cb"]) def test_causal_lm_unsupported(cb): model = AutoModelForCausalLM.from_config(AutoConfig.for_model("opt")) + qeff_model = QEFFAutoModelForCausalLM(model, cb) with pytest.warns(): - QEFFAutoModelForCausalLM(model, cb) + qeff_model.transform() @pytest.mark.parametrize("cb", [False, True], ids=["nocb", "cb"]) @@ -94,6 +95,9 @@ def test_causal_lm_init(config, cb): qeff_model = QEFFAutoModelForCausalLM(model, cb) with pytest.raises(TypeError): QEFFAutoModelForCausalLM(AutoModel.from_config(config, **model_kwargs), cb) + assert not qeff_model.is_transformed + qeff_model.transform(seq_len=1, bs=1) + assert qeff_model.is_transformed assert qeff_model.model.__class__.__name__.startswith("QEff") @@ -104,6 +108,9 @@ def test_causal_lm_pretrained(config, cb, tmp_path): model.save_pretrained(tmp_path) qeff_model = QEFFAutoModelForCausalLM.from_pretrained(tmp_path, cb) + assert not qeff_model.is_transformed + qeff_model.transform(seq_len=1, bs=1) + assert qeff_model.is_transformed assert qeff_model.model.__class__.__name__.startswith("QEff") @@ -172,7 +179,7 @@ def test_causal_lm_hash_creation(config, cb, subfunc, prefill_only, tmp_path): model = AutoModelForCausalLM.from_config(config, **model_kwargs) qeff_model = QEFFAutoModelForCausalLM(model, cb) qeff_model.export(tmp_path, use_onnx_subfunctions=subfunc, prefill_only=prefill_only) - hash_params = {} + hash_params = copy.deepcopy(qeff_model.hash_params) hash_params["config"] = qeff_model.model.config.to_diff_dict() hash_params["peft_config"] = None hash_params["applied_transform_names"] = qeff_model._transform_names() diff --git a/tests/transformers/qeff_classes/test_automodel_for_speech_seq2seq.py b/tests/transformers/qeff_classes/test_automodel_for_speech_seq2seq.py index 990049d7ac..28869de538 100644 --- a/tests/transformers/qeff_classes/test_automodel_for_speech_seq2seq.py +++ b/tests/transformers/qeff_classes/test_automodel_for_speech_seq2seq.py @@ -55,8 +55,9 @@ def test_seq2seq_unsupported(): model = AutoModelForSpeechSeq2Seq.from_config(AutoConfig.for_model("speech_to_text")) + qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) with pytest.warns(): - QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform() @pytest.mark.parametrize("config", configs, ids=config_ids) @@ -65,6 +66,9 @@ def test_seq2seq_init(config): qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) with pytest.raises(TypeError): QEFFAutoModelForSpeechSeq2Seq(AutoModel.from_config(config, **model_kwargs)) + assert not qeff_model.is_transformed + qeff_model.transform() + assert qeff_model.is_transformed assert qeff_model.model.model.__class__.__name__.startswith("QEff") @@ -74,6 +78,9 @@ def test_seq2seq_pretrained(config, tmp_path): model.save_pretrained(tmp_path) qeff_model = QEFFAutoModelForSpeechSeq2Seq.from_pretrained(tmp_path) + assert not qeff_model.is_transformed + qeff_model.transform() + assert qeff_model.is_transformed assert qeff_model.model.model.__class__.__name__.startswith("QEff") @@ -129,7 +136,7 @@ def test_seq2seq_hash_creation(config, tmp_path): model = AutoModelForSpeechSeq2Seq.from_config(config, **model_kwargs) qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) qeff_model.export(tmp_path) - hash_params = {} + hash_params = copy.deepcopy(qeff_model.hash_params) hash_params["config"] = qeff_model.model.config.to_diff_dict() hash_params["peft_config"] = None hash_params["applied_transform_names"] = qeff_model._transform_names() From 6253af1ad78256d27948dae65832f4b82b2b2472 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 7 Sep 2026 12:05:54 +0530 Subject: [PATCH 2/3] fix tests Signed-off-by: Mamta Singh --- tests/dynamo/test_transforms.py | 2 + tests/transformers/test_pytorch_transforms.py | 5 +- tests/unit_test/e2e/test_speech_e2e.py | 5 ++ tests/unit_test/e2e/test_vlm_e2e.py | 17 ++++--- .../models/test_causal_lm_accuracy.py | 30 +++++++----- .../unit_test/models/test_gemma2_accuracy.py | 46 +++++++++++-------- .../unit_test/models/test_gemma4_accuracy.py | 4 +- .../unit_test/models/test_model_quickcheck.py | 14 ++++-- .../models/test_prefill_decode_kv_handoff.py | 16 +++++-- .../transforms/test_peft_transforms.py | 1 + .../test_quantization_transforms.py | 12 +++-- .../transforms/test_transform_accuracy.py | 13 +++--- tests/unit_test/utils/test_auto_model_api.py | 9 +++- tests/unit_test/utils/test_error_handling.py | 3 +- tests/unit_test/utils/test_input_handler.py | 1 + .../unit_test/utils/test_modeling_registry.py | 18 ++++++-- tests/weight_free/test_transforms.py | 2 + 17 files changed, 132 insertions(+), 66 deletions(-) diff --git a/tests/dynamo/test_transforms.py b/tests/dynamo/test_transforms.py index 052ab0a098..fa14560635 100644 --- a/tests/dynamo/test_transforms.py +++ b/tests/dynamo/test_transforms.py @@ -132,6 +132,7 @@ class TestTemporarilyEnableNestedCompileRegions: def test_patches_decoder_layers_and_restores(self): model_hf, _ = make_tiny_llama() qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform() inner_model = qeff_model.model decoder_layers = [m for m in inner_model.modules() if isinstance(m, QEffLlamaDecoderLayer)] @@ -159,6 +160,7 @@ def test_patches_decoder_layers_and_restores(self): def test_noop_when_already_wrapped(self): model_hf, _ = make_tiny_llama() qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform() inner_model = qeff_model.model decoder_layers = [m for m in inner_model.modules() if isinstance(m, QEffLlamaDecoderLayer)] diff --git a/tests/transformers/test_pytorch_transforms.py b/tests/transformers/test_pytorch_transforms.py index 1e7dfd1088..442609e7ad 100644 --- a/tests/transformers/test_pytorch_transforms.py +++ b/tests/transformers/test_pytorch_transforms.py @@ -190,7 +190,10 @@ def run_kv_cache_transform_and_test( qaic_config = None if "num_logits_to_keep" in qaic_model_inputs: qaic_config = dict(speculative_model_type="target") - hf_model = QEFFAutoModelForCausalLM(hf_model, qaic_config=qaic_config).model + qeff_model = QEFFAutoModelForCausalLM(hf_model, qaic_config=qaic_config) + ctx_len = qaic_model_inputs["past_key_values"][0][0].shape[2] + qeff_model.transform(ctx_len=ctx_len, seq_len=input_len, bs=input_ids.shape[0], qaic_config=qaic_config) + hf_model = qeff_model.model if hidden_size_projections is not None: hf_model.projections = hidden_size_projections diff --git a/tests/unit_test/e2e/test_speech_e2e.py b/tests/unit_test/e2e/test_speech_e2e.py index 71f9b50c57..278345bb53 100644 --- a/tests/unit_test/e2e/test_speech_e2e.py +++ b/tests/unit_test/e2e/test_speech_e2e.py @@ -125,6 +125,7 @@ def test_qeff_whisper_model_class_replaced(self): model, cfg = make_tiny_whisper() qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform(ctx_len=MAX_TARGET_POS, seq_len=1, bs=1) assert isinstance(qeff_model.model, QEffWhisperForConditionalGeneration), ( f"Expected QEffWhisperForConditionalGeneration, got {type(qeff_model.model)}" ) @@ -134,6 +135,7 @@ def test_qeff_whisper_encoder_replaced(self): model, cfg = make_tiny_whisper() qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform(ctx_len=MAX_TARGET_POS, seq_len=1, bs=1) assert isinstance(qeff_model.model.model.encoder, QEffWhisperEncoder), ( f"Expected QEffWhisperEncoder, got {type(qeff_model.model.model.encoder)}" ) @@ -143,6 +145,7 @@ def test_qeff_whisper_decoder_replaced(self): model, cfg = make_tiny_whisper() qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform(ctx_len=MAX_TARGET_POS, seq_len=1, bs=1) assert isinstance(qeff_model.model.model.decoder, QEffWhisperDecoder), ( f"Expected QEffWhisperDecoder, got {type(qeff_model.model.model.decoder)}" ) @@ -152,6 +155,7 @@ def test_qeff_whisper_has_qeff_attention_layers(self): model, cfg = make_tiny_whisper() qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform(ctx_len=MAX_TARGET_POS, seq_len=1, bs=1) has_qeff_attn = any(isinstance(m, QEffWhisperAttention) for m in qeff_model.model.modules()) assert has_qeff_attn, "QEff Whisper must have QEffWhisperAttention layers" @@ -160,6 +164,7 @@ def test_qeff_whisper_has_positional_embedding_replaced(self): model, cfg = make_tiny_whisper() qeff_model = QEFFAutoModelForSpeechSeq2Seq(model) + qeff_model.transform(ctx_len=MAX_TARGET_POS, seq_len=1, bs=1) has_pos_emb = any(isinstance(m, QEffWhisperPositionalEmbedding) for m in qeff_model.model.modules()) assert has_pos_emb, "QEff Whisper must have QEffWhisperPositionalEmbedding" diff --git a/tests/unit_test/e2e/test_vlm_e2e.py b/tests/unit_test/e2e/test_vlm_e2e.py index 52ba976613..5ee8806e71 100644 --- a/tests/unit_test/e2e/test_vlm_e2e.py +++ b/tests/unit_test/e2e/test_vlm_e2e.py @@ -20,6 +20,11 @@ import pytest + +def get_pytorch_transform_pipeline(wrapper_cls): + return wrapper_cls.__new__(wrapper_cls)._all_pytorch_transforms() + + # --------------------------------------------------------------------------- # Tests: QEFFAutoModelForImageTextToText class structure # --------------------------------------------------------------------------- @@ -315,8 +320,8 @@ def test_pytorch_transforms_include_custom_ops_transform(self): from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCTC from QEfficient.transformers.models.pytorch_transforms import CustomOpsTransform - assert CustomOpsTransform in QEFFAutoModelForCTC._pytorch_transforms, ( - "CustomOpsTransform not in QEFFAutoModelForCTC._pytorch_transforms" + assert CustomOpsTransform in get_pytorch_transform_pipeline(QEFFAutoModelForCTC), ( + "CustomOpsTransform not in QEFFAutoModelForCTC resolved pytorch transforms" ) def test_onnx_transforms_include_fp16_clip(self): @@ -372,8 +377,8 @@ def test_single_qpc_pytorch_transforms_include_kv_offload_transform(self): from QEfficient.transformers.models.modeling_auto import _QEFFAutoModelForImageTextToTextSingleQPC from QEfficient.transformers.models.pytorch_transforms import VlmNoKVOffloadTransform - assert VlmNoKVOffloadTransform in _QEFFAutoModelForImageTextToTextSingleQPC._pytorch_transforms, ( - "VlmNoKVOffloadTransform not in SingleQPC._pytorch_transforms" + assert VlmNoKVOffloadTransform in get_pytorch_transform_pipeline(_QEFFAutoModelForImageTextToTextSingleQPC), ( + "VlmNoKVOffloadTransform not in SingleQPC resolved pytorch transforms" ) def test_single_qpc_pytorch_transforms_include_no_kv_offload(self): @@ -381,6 +386,6 @@ def test_single_qpc_pytorch_transforms_include_no_kv_offload(self): from QEfficient.transformers.models.modeling_auto import _QEFFAutoModelForImageTextToTextSingleQPC from QEfficient.transformers.models.pytorch_transforms import VlmNoKVOffloadTransform - assert VlmNoKVOffloadTransform in _QEFFAutoModelForImageTextToTextSingleQPC._pytorch_transforms, ( - "VlmNoKVOffloadTransform not in SingleQPC._pytorch_transforms" + assert VlmNoKVOffloadTransform in get_pytorch_transform_pipeline(_QEFFAutoModelForImageTextToTextSingleQPC), ( + "VlmNoKVOffloadTransform not in SingleQPC resolved pytorch transforms" ) diff --git a/tests/unit_test/models/test_causal_lm_accuracy.py b/tests/unit_test/models/test_causal_lm_accuracy.py index ccf455a3c6..1c5e66819e 100644 --- a/tests/unit_test/models/test_causal_lm_accuracy.py +++ b/tests/unit_test/models/test_causal_lm_accuracy.py @@ -82,6 +82,12 @@ def make_qeff_inputs(input_ids, config, ctx_len=CTX_LEN): return {"input_ids": input_ids, "position_ids": position_ids, "past_key_values": past_key_values} +def make_transformed_qeff_model(model, continuous_batching=False): + qeff_model = QEFFAutoModelForCausalLM(model, continuous_batching=continuous_batching) + qeff_model.transform(ctx_len=CTX_LEN, seq_len=SEQ_LEN, bs=1) + return qeff_model + + # --------------------------------------------------------------------------- # Tiny model factories # --------------------------------------------------------------------------- @@ -257,7 +263,7 @@ def _assert_same_greedy_token(self, model, cfg, label): hf_logits = model(input_ids=input_ids).logits[:, -1, :] hf_token = hf_logits.argmax(-1).item() - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): qeff_logits = qeff_model.model(**qeff_inputs).logits[:, -1, :] @@ -274,7 +280,7 @@ def _assert_logits_numerically_close(self, model, cfg, label, atol=1e-3): with torch.no_grad(): hf_logits = model(input_ids=input_ids).logits[:, -1, :] - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): qeff_logits = qeff_model.model(**qeff_inputs).logits[:, -1, :] @@ -338,7 +344,7 @@ def test_qeff_logits_are_finite(self): (make_tiny_phi3, "Phi3"), ]: model, cfg = factory() - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN)) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): @@ -348,7 +354,7 @@ def test_qeff_logits_are_finite(self): def test_qeff_past_key_values_returned(self): """QEff model must return past_key_values for the decode step.""" model, cfg = make_tiny_gpt2() - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN)) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): @@ -363,7 +369,7 @@ def test_gpt2_top5_tokens_overlap_with_hf(self): with torch.no_grad(): hf_top5 = set(model(input_ids=input_ids).logits[:, -1, :].topk(5).indices.squeeze().tolist()) - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): qeff_top5 = set(qeff_model.model(**qeff_inputs).logits[:, -1, :].topk(5).indices.squeeze().tolist()) @@ -384,7 +390,7 @@ class TestQEffDecodeStepAccuracy: def _run_prefill_then_decode(self, model, cfg, n_decode_steps=3, input_ids=None): """Run prefill + n decode steps, return list of generated token IDs.""" - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) if input_ids is None: input_ids = torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN)) qeff_inputs = make_qeff_inputs(input_ids, cfg) @@ -448,7 +454,7 @@ def test_gpt2_prefill_token_matches_hf_next_token(self): with torch.no_grad(): hf_next = model(input_ids=input_ids).logits[:, -1, :].argmax(-1).item() - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): qeff_next = qeff_model.model(**qeff_inputs).logits[:, -1, :].argmax(-1).item() @@ -462,7 +468,7 @@ def test_llama_prefill_token_matches_hf_next_token(self): with torch.no_grad(): hf_next = model(input_ids=input_ids).logits[:, -1, :].argmax(-1).item() - qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model = make_transformed_qeff_model(model) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): qeff_next = qeff_model.model(**qeff_inputs).logits[:, -1, :].argmax(-1).item() @@ -495,13 +501,13 @@ class TestContinuousBatchingMode: def test_gpt2_continuous_batching_wraps_without_error(self): model, cfg = make_tiny_gpt2() - qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + qeff = make_transformed_qeff_model(model, continuous_batching=True) assert qeff is not None assert qeff.continuous_batching is True def test_llama_continuous_batching_wraps_without_error(self): model, cfg = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + qeff = make_transformed_qeff_model(model, continuous_batching=True) assert qeff is not None assert qeff.continuous_batching is True @@ -510,7 +516,7 @@ def test_gpt2_continuous_batching_model_is_transformed(self): from QEfficient.transformers.models.gpt2.modeling_gpt2 import QEffGPT2LMHeadModel model, cfg = make_tiny_gpt2() - qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + qeff = make_transformed_qeff_model(model, continuous_batching=True) assert isinstance(qeff.model, QEffGPT2LMHeadModel) def test_continuous_batching_false_is_default(self): @@ -521,7 +527,7 @@ def test_continuous_batching_false_is_default(self): def test_continuous_batching_model_produces_finite_logits(self): """Continuous batching model must produce finite logits.""" model, cfg = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + qeff = make_transformed_qeff_model(model, continuous_batching=True) input_ids = torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN)) qeff_inputs = make_qeff_inputs(input_ids, cfg) with torch.no_grad(): diff --git a/tests/unit_test/models/test_gemma2_accuracy.py b/tests/unit_test/models/test_gemma2_accuracy.py index 2d3527fa5d..c698c3622b 100644 --- a/tests/unit_test/models/test_gemma2_accuracy.py +++ b/tests/unit_test/models/test_gemma2_accuracy.py @@ -106,6 +106,12 @@ def _decode_inputs(next_token, decode_position, past_key_values): } +def _make_transformed_qeff_gemma2(model): + qeff = QEFFAutoModelForCausalLM(model) + qeff.transform(ctx_len=CTX_LEN, seq_len=PREFILL_LEN, bs=1) + return qeff + + def _extract_next_token(logits): """ Extract greedy next token. QEffGemma2ForCausalLM returns (batch, 1, vocab), @@ -166,24 +172,24 @@ class TestQEffGemma2Architecture: def test_qeff_wraps_without_error(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) assert qeff is not None assert hasattr(qeff, "model") def test_qeff_model_class_is_qeff_gemma2(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) assert isinstance(qeff.model, QEffGemma2ForCausalLM), f"Expected QEffGemma2ForCausalLM, got {type(qeff.model)}" def test_qeff_model_is_eval_mode(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) assert not qeff.model.training def test_qeff_model_has_same_parameter_count_as_hf(self): model, cfg = make_tiny_gemma2() hf_params = sum(p.numel() for p in model.parameters()) - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) qeff_params = sum(p.numel() for p in qeff.model.parameters()) # QEffGemma2Model registers sin_cached and cos_cached as nn.Parameter, # which adds extra parameters compared to the HF model. Allow for this. @@ -210,7 +216,7 @@ def test_prefill_logits_shape_is_batch_1_vocab(self): not (1, PREFILL_LEN, VOCAB_SIZE). """ model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -223,7 +229,7 @@ def test_prefill_logits_shape_is_batch_1_vocab(self): def test_decode_logits_shape_is_batch_1_vocab(self): """QEff Gemma2 decode must also return (1, 1, VOCAB_SIZE).""" model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): prefill_out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -236,7 +242,7 @@ def test_decode_logits_shape_is_batch_1_vocab(self): def test_prefill_logits_are_finite(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -264,7 +270,7 @@ def test_prefill_token_matches_hf(self): with torch.no_grad(): hf_token = model(input_ids=input_ids).logits[:, -1, :].argmax(-1).item() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) with torch.no_grad(): qeff_out = qeff.model(**_prefill_inputs(input_ids, cfg)) qeff_token = _extract_next_token(qeff_out.logits) @@ -282,7 +288,7 @@ def test_prefill_logits_numerically_close_to_hf(self): with torch.no_grad(): hf_logits = model(input_ids=input_ids).logits[:, -1, :] - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) with torch.no_grad(): qeff_out = qeff.model(**_prefill_inputs(input_ids, cfg)) # qeff_out.logits is (1, 1, vocab) — squeeze to (1, vocab) @@ -301,7 +307,7 @@ def test_top5_tokens_overlap_with_hf(self): with torch.no_grad(): hf_top5 = set(model(input_ids=input_ids).logits[:, -1, :].topk(5).indices.squeeze().tolist()) - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) with torch.no_grad(): qeff_out = qeff.model(**_prefill_inputs(input_ids, cfg)) qeff_top5 = set(qeff_out.logits[:, -1, :].topk(5).indices.squeeze().tolist()) @@ -326,7 +332,7 @@ class TestQEffGemma2CacheWritten: def test_past_key_values_not_none_after_prefill(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -338,7 +344,7 @@ def test_cache_is_non_zero_after_prefill(self): At least one position in the prefill range must be non-zero. """ model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -365,7 +371,7 @@ def test_cache_is_non_zero_after_prefill(self): def test_cache_has_correct_number_of_layers(self): """past_key_values must have one entry per transformer layer.""" model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -399,7 +405,7 @@ class TestQEffGemma2PrefillDecodeHandoff: def test_decode_with_real_cache_produces_valid_token(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): @@ -414,7 +420,7 @@ def test_decode_with_real_cache_produces_valid_token(self): def test_decode_with_real_cache_returns_finite_logits(self): model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): @@ -429,7 +435,7 @@ def test_decode_with_real_cache_returns_finite_logits(self): def test_three_decode_steps_all_valid(self): """Three consecutive decode steps with real cache must all produce valid tokens.""" model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): @@ -455,7 +461,7 @@ def test_three_decode_steps_all_valid(self): def test_three_decode_steps_all_finite(self): """All decode logits must be finite.""" model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): @@ -482,7 +488,7 @@ def test_decode_is_deterministic(self): input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) def _run(m): - qeff = QEFFAutoModelForCausalLM(m) + qeff = _make_transformed_qeff_gemma2(m) with torch.no_grad(): prefill_out = qeff.model(**_prefill_inputs(input_ids, cfg)) token = _extract_next_token(prefill_out.logits) @@ -510,7 +516,7 @@ def test_real_cache_differs_from_zero_cache(self): for seed in range(8): torch.manual_seed(seed) - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): @@ -540,7 +546,7 @@ def test_real_cache_differs_from_zero_cache(self): def test_decode_position_advances_strictly(self): """Each decode step must use a strictly increasing position_id.""" model, cfg = make_tiny_gemma2() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_gemma2(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): diff --git a/tests/unit_test/models/test_gemma4_accuracy.py b/tests/unit_test/models/test_gemma4_accuracy.py index 026342cce8..e7d9c41329 100644 --- a/tests/unit_test/models/test_gemma4_accuracy.py +++ b/tests/unit_test/models/test_gemma4_accuracy.py @@ -162,7 +162,9 @@ def _make_qeff_gemma4(model): """ Use the ImageTextToText auto-wrapper for this Gemma4 test model. """ - return QEFFAutoModelForImageTextToText(model) + qeff = QEFFAutoModelForImageTextToText(model) + qeff.transform(ctx_len=CTX_LEN, seq_len=PREFILL_LEN, bs=1) + return qeff def _qeff_forward(qeff, **inputs): diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 5fe9217794..83670e0566 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -1036,6 +1036,8 @@ def test_kimi_k25_quickcheck_hf_qeff_vision_logits_parity(): config=model_hf.config, torch_dtype=torch.float32, ) + seq_len = inputs["input_ids"].shape[1] + qeff_model.transform(ctx_len=seq_len, seq_len=seq_len, bs=inputs["input_ids"].shape[0]) grid_thws = inputs["grid_thws"].to(torch.int64) h_shape = torch.ones((int(grid_thws[0, 1].item()),), dtype=torch.int64) w_shape = torch.ones((int(grid_thws[0, 2].item()),), dtype=torch.int64) @@ -1084,6 +1086,7 @@ def test_causal_lm_cpu_runtime_parity_with_api_runner(model_type, model_id, tmp_ hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform(ctx_len=ctx_len, seq_len=prompt_len, bs=1) kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) onnx_path = _exported_onnx_path(qeff_model.export(tmp_path)) ort_tokens = api_runner.run_kv_model_on_ort(str(onnx_path)) @@ -1113,6 +1116,7 @@ def test_vlm_text_side_runtime_parity_and_full_export(tmp_path): hf_tokens = api_runner.run_hf_model_on_pytorch(text_model) qeff_text_model = QEFFAutoModelForCausalLM(text_model) + qeff_text_model.transform(ctx_len=8, seq_len=4, bs=1) kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_text_model.model) onnx_path = _exported_onnx_path(qeff_text_model.export(tmp_path / "vlm-text")) ort_tokens = api_runner.run_kv_model_on_ort(str(onnx_path)) @@ -1500,6 +1504,7 @@ def test_causal_subfunction_export_smoke_all_models(model_type, model_id, tmp_pa hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform(ctx_len=ctx_len, seq_len=prompt_len, bs=1) kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) onnx_path = _exported_onnx_path(qeff_model.export(tmp_path / "with-subfunctions-all", use_onnx_subfunctions=True)) ort_tokens = api_runner.run_kv_model_on_ort(str(onnx_path)) @@ -2749,12 +2754,13 @@ def test_layerwise_matches_default_path_for_qwen3_moe(): torch.manual_seed(0) hf = Qwen3MoeForCausalLM(cfg).eval() qeff_model = QEfficient.QEFFAutoModelForCausalLM(hf, continuous_batching=False) - inner = qeff_model.model.model B, S, ctx, num_layers = 1, 8, 16, cfg.num_hidden_layers n_kv, head_dim = cfg.num_key_value_heads, cfg.head_dim ids = torch.randint(0, cfg.vocab_size, (B, S)) position_ids = torch.arange(S).view(1, -1) + qeff_model.transform(ctx_len=ctx, seq_len=S, bs=B) + inner = qeff_model.model.model def fresh_pkv(): return tuple( @@ -2815,10 +2821,11 @@ def test_layerwise_matches_default_path_for_qwen3_5_moe(): dtype=torch.float32, layerwise=False, ) - wrapper = qeff_model.model.get_qeff_language_decoder().eval() + qeff_model.transform(ctx_len=16, seq_len=8, bs=1) lang_inputs = qeff_model.model.get_dummy_inputs(kv_offload=True)["lang"] lang_inputs["input_ids"][0, 0] = qeff_model.model.config.image_token_id lang_inputs["vision_embeds"].normal_() + wrapper = qeff_model.lang_model.model.eval() with torch.no_grad(): default_out = wrapper(**lang_inputs) @@ -2872,8 +2879,9 @@ def test_layerwise_matches_default_path_for_qwen3_vl_moe(): dtype=torch.float32, layerwise=False, ) - wrapper = qeff_model.model.get_qeff_language_decoder().eval() + qeff_model.transform(ctx_len=16, seq_len=8, bs=1) lang_inputs = qeff_model.model.get_dummy_inputs(kv_offload=True)["lang"] + wrapper = qeff_model.lang_model.model.eval() with torch.no_grad(): default_out = wrapper(**lang_inputs) diff --git a/tests/unit_test/models/test_prefill_decode_kv_handoff.py b/tests/unit_test/models/test_prefill_decode_kv_handoff.py index cd6b5cab65..41ea9885d5 100644 --- a/tests/unit_test/models/test_prefill_decode_kv_handoff.py +++ b/tests/unit_test/models/test_prefill_decode_kv_handoff.py @@ -111,6 +111,12 @@ def _decode_inputs(next_token, decode_position, past_key_values): } +def _make_transformed_qeff_model(model): + qeff = QEFFAutoModelForCausalLM(model) + qeff.transform(ctx_len=CTX_LEN, seq_len=PREFILL_LEN, bs=1) + return qeff + + # --------------------------------------------------------------------------- # Tiny model factories # --------------------------------------------------------------------------- @@ -212,7 +218,7 @@ def _run_real_handoff(factory, n_decode_steps=3, seed=42): """ torch.manual_seed(seed) model, cfg = factory() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) prefill_in = _prefill_inputs(input_ids, cfg) @@ -258,7 +264,7 @@ class TestPrefillWritesCache: def _assert_cache_written(self, factory, label): model, cfg = factory() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) with torch.no_grad(): out = qeff.model(**_prefill_inputs(input_ids, cfg)) @@ -390,7 +396,7 @@ def _assert_cache_influences_output(self, factory, label, n_seeds=8): for seed in range(n_seeds): torch.manual_seed(seed) - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) # Prefill to get real cache @@ -446,7 +452,7 @@ class TestDecodePositionAdvancesStrictly: def _assert_positions_advance(self, factory, label): model, cfg = factory() - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_model(model) input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN)) prefill_in = _prefill_inputs(input_ids, cfg) @@ -515,7 +521,7 @@ def _assert_full_pipeline(self, factory, label): hf_token = hf_logits.argmax(-1).item() # QEff prefill - qeff = QEFFAutoModelForCausalLM(model) + qeff = _make_transformed_qeff_model(model) with torch.no_grad(): prefill_out = qeff.model(**_prefill_inputs(input_ids, cfg)) qeff_token = _extract_next_token(prefill_out.logits) diff --git a/tests/unit_test/transforms/test_peft_transforms.py b/tests/unit_test/transforms/test_peft_transforms.py index 80c1dcf469..341289c8b3 100644 --- a/tests/unit_test/transforms/test_peft_transforms.py +++ b/tests/unit_test/transforms/test_peft_transforms.py @@ -203,6 +203,7 @@ def test_qeff_peft_model_forward_produces_finite_logits(self): lora_model, cfg = self._make_lora_model() qeff_peft = QEffAutoPeftModelForCausalLM(lora_model) + qeff_peft.transform(ctx_len=CTX_LEN, seq_len=SEQ_LEN, bs=1) n_layers = cfg.num_hidden_layers n_kv = cfg.num_key_value_heads diff --git a/tests/unit_test/transforms/test_quantization_transforms.py b/tests/unit_test/transforms/test_quantization_transforms.py index e89432873c..25723ce8a6 100644 --- a/tests/unit_test/transforms/test_quantization_transforms.py +++ b/tests/unit_test/transforms/test_quantization_transforms.py @@ -21,6 +21,10 @@ """ +def get_pytorch_transform_pipeline(wrapper_cls): + return wrapper_cls.__new__(wrapper_cls)._all_pytorch_transforms() + + # --------------------------------------------------------------------------- # Tests: Quantization Transform Importability and Structure # --------------------------------------------------------------------------- @@ -323,11 +327,11 @@ def test_quantization_transforms_come_before_kv_cache_transform(self): from QEfficient.transformers.models.pytorch_transforms import KVCacheTransform from QEfficient.transformers.quantizers.quant_transforms import AwqToMatmulNbitsTransform - transforms = QEFFAutoModelForCausalLM._pytorch_transforms + transforms = get_pytorch_transform_pipeline(QEFFAutoModelForCausalLM) awq_idx = next((i for i, t in enumerate(transforms) if t is AwqToMatmulNbitsTransform), None) kv_idx = next((i for i, t in enumerate(transforms) if t is KVCacheTransform), None) - assert awq_idx is not None, "AwqToMatmulNbitsTransform not found in _pytorch_transforms" - assert kv_idx is not None, "KVCacheTransform not found in _pytorch_transforms" + assert awq_idx is not None, "AwqToMatmulNbitsTransform not found in resolved pytorch transforms" + assert kv_idx is not None, "KVCacheTransform not found in resolved pytorch transforms" assert awq_idx < kv_idx, ( f"AwqToMatmulNbitsTransform (idx={awq_idx}) must come before KVCacheTransform (idx={kv_idx})" ) @@ -341,7 +345,7 @@ def test_image_text_wrappers_include_pack_quantized_int4_transform(self): from QEfficient.transformers.quantizers.quant_transforms import PackQuantizedInt4ToMatMulNBitsTransform for wrapper_cls in [QEffVisionEncoderForTextImageToTextModel, QEffCausalLMForTextImageToTextModel]: - transforms = wrapper_cls._pytorch_transforms + transforms = get_pytorch_transform_pipeline(wrapper_cls) pack_idx = transforms.index(PackQuantizedInt4ToMatMulNBitsTransform) custom_ops_idx = transforms.index(CustomOpsTransform) assert pack_idx < custom_ops_idx diff --git a/tests/unit_test/transforms/test_transform_accuracy.py b/tests/unit_test/transforms/test_transform_accuracy.py index 3b6657f8fd..fe061b42f0 100644 --- a/tests/unit_test/transforms/test_transform_accuracy.py +++ b/tests/unit_test/transforms/test_transform_accuracy.py @@ -87,6 +87,10 @@ CTX_LEN = 32 +def get_pytorch_transform_pipeline(wrapper_cls): + return wrapper_cls.__new__(wrapper_cls)._all_pytorch_transforms() + + # --------------------------------------------------------------------------- # Tiny model factories # --------------------------------------------------------------------------- @@ -407,8 +411,8 @@ def test_repeat_kv_skips_encoder_wrapper_without_config(self): ) model_hf = AutoModelForImageTextToText.from_config(cfg) qeff_model = QEFFAutoModelForImageTextToText(copy.deepcopy(model_hf), kv_offload=True, qaic_config={}) - assert not hasattr(qeff_model.vision_model.model, "config") qeff_model.vision_model.transform(ctx_len=64, seq_len=8, bs=1, qaic_config={"replicate_kv_heads": True}) + assert not hasattr(qeff_model.vision_model.model, "config") assert qeff_model.vision_model.hash_params["num_replicate_kv_heads"] == 1 def test_calculate_num_replicate_kv_heads_for_gqa_mqa_and_mha(self): @@ -1404,16 +1408,13 @@ def test_simple_decode_moe_transform_is_registered_after_cache_transforms(self): _QEFFAutoModelForImageTextToTextSingleQPC, ) - assert QEFFAutoModelForCausalLM._pytorch_transforms[-1] is SimpleDecodeMoeTransform - assert QEffCausalLMForTextImageToTextModel._pytorch_transforms[-1] is SimpleDecodeMoeTransform - assert _QEFFAutoModelForImageTextToTextSingleQPC._pytorch_transforms[-1] is SimpleDecodeMoeTransform - for wrapper in ( QEFFAutoModelForCausalLM, QEffCausalLMForTextImageToTextModel, _QEFFAutoModelForImageTextToTextSingleQPC, ): - transforms = wrapper._pytorch_transforms + transforms = get_pytorch_transform_pipeline(wrapper) + assert transforms[-1] is SimpleDecodeMoeTransform assert transforms.index(KVCacheTransform) < transforms.index(SimpleDecodeMoeTransform) def test_moe_component_mappings_owned_by_optimized_mapper(self): diff --git a/tests/unit_test/utils/test_auto_model_api.py b/tests/unit_test/utils/test_auto_model_api.py index 07428e94e7..85d799ad46 100644 --- a/tests/unit_test/utils/test_auto_model_api.py +++ b/tests/unit_test/utils/test_auto_model_api.py @@ -24,6 +24,11 @@ import torch from transformers import GPT2Config, GPT2LMHeadModel + +def get_pytorch_transform_pipeline(wrapper_cls): + return wrapper_cls.__new__(wrapper_cls)._all_pytorch_transforms() + + # --------------------------------------------------------------------------- # Tiny model factories # --------------------------------------------------------------------------- @@ -417,14 +422,14 @@ def test_pytorch_transforms_contains_kv_cache_transform(self): from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM from QEfficient.transformers.models.pytorch_transforms import KVCacheTransform - assert KVCacheTransform in QEFFAutoModelForCausalLM._pytorch_transforms + assert KVCacheTransform in get_pytorch_transform_pipeline(QEFFAutoModelForCausalLM) def test_pytorch_transforms_contains_custom_ops_transform(self): """_pytorch_transforms must contain CustomOpsTransform.""" from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM from QEfficient.transformers.models.pytorch_transforms import CustomOpsTransform - assert CustomOpsTransform in QEFFAutoModelForCausalLM._pytorch_transforms + assert CustomOpsTransform in get_pytorch_transform_pipeline(QEFFAutoModelForCausalLM) def test_has_onnx_transforms_list(self): """QEFFAutoModelForCausalLM must have _onnx_transforms list.""" diff --git a/tests/unit_test/utils/test_error_handling.py b/tests/unit_test/utils/test_error_handling.py index c0fb7da665..7663457e2a 100644 --- a/tests/unit_test/utils/test_error_handling.py +++ b/tests/unit_test/utils/test_error_handling.py @@ -344,8 +344,9 @@ def test_is_tlm_true_with_target_type(self): def test_turbo_type_requires_pretrained_model_name(self): """speculative_model_type='turbo' without pretrained_model_name_or_path must raise KeyError.""" model = make_tiny_llama() + qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "turbo"}) with pytest.raises(KeyError, match="pretrained_model_name_or_path"): - QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "turbo"}) + qeff.transform() def test_cb_and_tlm_together_model_is_tlm(self): """continuous_batching=True with TLM: model must still be recognized as TLM.""" diff --git a/tests/unit_test/utils/test_input_handler.py b/tests/unit_test/utils/test_input_handler.py index ef964529ba..b9c8fc5121 100644 --- a/tests/unit_test/utils/test_input_handler.py +++ b/tests/unit_test/utils/test_input_handler.py @@ -175,6 +175,7 @@ def _run_prefill(self, tok, cfg, prompt_len=8): model = GPT2LMHeadModel(cfg).eval() qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model.transform() handler = _make_handler(tok, cfg, prompt_len=prompt_len) inputs = handler.prepare_pytorch_inputs() with torch.no_grad(): diff --git a/tests/unit_test/utils/test_modeling_registry.py b/tests/unit_test/utils/test_modeling_registry.py index a08a52bea1..d2d39f447d 100644 --- a/tests/unit_test/utils/test_modeling_registry.py +++ b/tests/unit_test/utils/test_modeling_registry.py @@ -36,6 +36,11 @@ QEFFAutoModelForSpeechSeq2Seq, ) + +def get_pytorch_transform_pipeline(wrapper_cls): + return wrapper_cls.__new__(wrapper_cls)._all_pytorch_transforms() + + # --------------------------------------------------------------------------- # Tests: qeff_supported_architectures # --------------------------------------------------------------------------- @@ -423,18 +428,20 @@ def test_has_onnx_transforms_list(self): def test_kv_cache_transform_in_pytorch_transforms(self): transform_names = [ - t.__name__ if hasattr(t, "__name__") else str(t) for t in QEFFAutoModelForCausalLM._pytorch_transforms + t.__name__ if hasattr(t, "__name__") else str(t) + for t in get_pytorch_transform_pipeline(QEFFAutoModelForCausalLM) ] assert any("KVCache" in name for name in transform_names), ( - f"KVCacheTransform not found in _pytorch_transforms: {transform_names}" + f"KVCacheTransform not found in resolved pytorch transforms: {transform_names}" ) def test_custom_ops_transform_in_pytorch_transforms(self): transform_names = [ - t.__name__ if hasattr(t, "__name__") else str(t) for t in QEFFAutoModelForCausalLM._pytorch_transforms + t.__name__ if hasattr(t, "__name__") else str(t) + for t in get_pytorch_transform_pipeline(QEFFAutoModelForCausalLM) ] assert any("CustomOps" in name for name in transform_names), ( - f"CustomOpsTransform not found in _pytorch_transforms: {transform_names}" + f"CustomOpsTransform not found in resolved pytorch transforms: {transform_names}" ) def test_has_hf_auto_class(self): @@ -477,7 +484,7 @@ def test_model_name_property_returns_string(self): assert len(qeff.model_name) > 0 def test_model_attribute_is_transformed_model(self): - """After construction, qeff.model must be the KV-transformed model.""" + """After transform(), qeff.model must be the KV-transformed model.""" from transformers import GPT2Config, GPT2LMHeadModel from QEfficient.transformers.models.gpt2.modeling_gpt2 import QEffGPT2LMHeadModel @@ -485,6 +492,7 @@ def test_model_attribute_is_transformed_model(self): cfg = GPT2Config(n_layer=1, n_head=2, n_embd=64, vocab_size=500, n_positions=32, n_ctx=32) model = GPT2LMHeadModel(cfg) qeff = QEFFAutoModelForCausalLM(model) + qeff.transform() assert isinstance(qeff.model, QEffGPT2LMHeadModel), f"Expected QEffGPT2LMHeadModel, got {type(qeff.model)}" def test_onnx_transforms_contain_fp16_clip(self): diff --git a/tests/weight_free/test_transforms.py b/tests/weight_free/test_transforms.py index b813f2337b..c7bcf126fc 100644 --- a/tests/weight_free/test_transforms.py +++ b/tests/weight_free/test_transforms.py @@ -623,6 +623,7 @@ class TestTemporarilyEnableNestedCompileRegions: def test_patches_decoder_layers_and_restores(self): model_hf, _ = make_tiny_llama() qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform() inner_model = qeff_model.model decoder_layers = [m for m in inner_model.modules() if isinstance(m, QEffLlamaDecoderLayer)] @@ -649,6 +650,7 @@ def test_patches_decoder_layers_and_restores(self): def test_noop_when_already_wrapped(self): model_hf, _ = make_tiny_llama() qeff_model = QEFFAutoModelForCausalLM(model_hf) + qeff_model.transform() inner_model = qeff_model.model decoder_layers = [m for m in inner_model.modules() if isinstance(m, QEffLlamaDecoderLayer)] From ed95a95d9d98d20383d6248ae979fbae9cdd2e64 Mon Sep 17 00:00:00 2001 From: Mamta Singh Date: Mon, 7 Sep 2026 15:51:54 +0530 Subject: [PATCH 3/3] accept qaic_config only via compile and deprecate export as external user facing API Signed-off-by: Mamta Singh --- QEfficient/__init__.py | 3 +- QEfficient/base/common.py | 16 ++ QEfficient/base/modeling_qeff.py | 48 +++- QEfficient/cloud/infer.py | 2 +- .../exporter/export_hf_to_cloud_ai_100.py | 11 +- .../generation/text_generation_inference.py | 9 +- .../transformers/models/modeling_auto.py | 192 +++++++------- QEfficient/utils/export_utils.py | 2 +- QEfficient/utils/test_utils.py | 9 +- docs/source/features_enablement.md | 8 +- docs/source/qeff_autoclasses.md | 7 - docs/source/quick_start.md | 8 +- .../embeddings/qwen3vl/qwen3_vl_embedding.py | 2 +- .../kimi_k2/example_kimi_k25_vision_disagg.py | 1 - .../models/kimi_k2/export_kimik2.py | 106 ++++---- .../models/qwen3vl/qwen3_vl_blocked.py | 3 +- .../compute_context_length/basic_inference.py | 6 +- .../compute_context_length/gemma3.py | 10 +- .../compute_context_length/gpt_oss.py | 9 +- .../gpt_oss_disagg_mode_with_chunking.py | 12 +- .../compute_context_length/granite_vision.py | 8 +- .../compute_context_length/internvl.py | 8 +- .../compute_context_length/llama4.py | 10 +- .../compute_context_length/llama4_cb.py | 235 +++++++++--------- .../llama4_multi_image.py | 9 +- .../compute_context_length/mistral3.py | 8 +- .../compute_context_length/molmo.py | 10 +- .../compute_context_length/qwen2_5_vl.py | 10 +- .../compute_context_length/qwen2_5_vl_cb.py | 9 +- .../compute_context_length/qwen3moe.py | 9 +- .../ccl_qwen3moe_inference.py | 7 +- .../compute_context_length/vlm_inference.py | 8 +- examples/performance/on_device_sampling.py | 3 +- .../speculative_decoding/draft_based.py | 6 +- .../speculative_decoding/multi_projection.py | 2 +- .../speculative_decoding/prompt_lookup.py | 6 +- examples/text_generation/run_kimik2.py | 2 +- .../causal_lm_models/check_causal_models.py | 4 +- .../causal_lm_models/test_fp16_causal_lm.py | 5 +- .../test_qwen3vl_embedding_mad.py | 2 +- .../test_image_text_to_text_models.py | 3 - .../test_causal_lm_blocking_subfunction.py | 8 +- tests/transformers/test_pytorch_transforms.py | 2 +- .../unit_test/models/test_model_quickcheck.py | 5 +- .../models/test_modeling_auto_cpu.py | 84 ++++++- .../transforms/test_speculative_decoding.py | 12 +- .../transforms/test_transform_accuracy.py | 6 +- tests/unit_test/utils/test_error_handling.py | 36 +-- tests/weight_free/test_ccl.py | 6 +- 49 files changed, 569 insertions(+), 418 deletions(-) diff --git a/QEfficient/__init__.py b/QEfficient/__init__.py index fb85183144..1f3f53809b 100755 --- a/QEfficient/__init__.py +++ b/QEfficient/__init__.py @@ -73,11 +73,10 @@ class HybridChunkedCache(HybridCache): warnings.formatwarning = custom_format_warning -# Users can use QEfficient.export for exporting models to ONNX +# Backward-compatible deprecated alias for exporting models to ONNX. Prefer .compile(). export = qualcomm_efficient_converter __all__ = [ "transform", - "export", "compile", "cloud_ai_100_exec_kv", "QEFFAutoModel", diff --git a/QEfficient/base/common.py b/QEfficient/base/common.py index 7f66b9f3f3..26a9b300ac 100644 --- a/QEfficient/base/common.py +++ b/QEfficient/base/common.py @@ -13,6 +13,7 @@ """ import os +import warnings from typing import Any from transformers import AutoConfig @@ -40,6 +41,14 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> """ Downloads HuggingFace model if already doesn't exist locally, returns QEFFAutoModel object based on type of model. """ + qaic_config = kwargs.pop("qaic_config", None) + if qaic_config is not None: + warnings.warn( + "Passing `qaic_config` to `from_pretrained()` is deprecated and will be removed in a future " + "release. Pass `qaic_config` to `compile()` instead.", + DeprecationWarning, + stacklevel=2, + ) config = AutoConfig.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) class_name = ( @@ -66,4 +75,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> continuous_batching=continuous_batching, **kwargs, ) + if qaic_config is not None: + if hasattr(qeff_model, "_activate_qaic_config"): + qeff_model._activate_qaic_config(qaic_config) + elif hasattr(qeff_model, "_resolve_qaic_config"): + qeff_model._resolve_qaic_config(qaic_config) + elif hasattr(qeff_model, "_set_qaic_config"): + qeff_model._set_qaic_config(qaic_config) return qeff_model diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 5f1c73800c..c9c53b162b 100644 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -5,6 +5,7 @@ # # ---------------------------------------------------------------------------- +import copy import gc import inspect import logging @@ -434,6 +435,7 @@ def maybe_apply_replicate_kv_transform(self, model_config, num_devices: int, qai def __init__(self, model: torch.nn.Module, **kwargs) -> None: super().__init__() + qaic_config = kwargs.pop("qaic_config", None) self.model = model self.config = model.config self.hash_params = create_model_params(self, **kwargs) @@ -452,6 +454,37 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None: self.is_transformed: bool = False self._normalize_torch_dtype() + self._set_qaic_config(qaic_config) + + def _copy_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + if qaic_config is None: + return None + if not isinstance(qaic_config, dict): + raise TypeError(f"`qaic_config` must be a dictionary, got {type(qaic_config).__name__}.") + return copy.deepcopy(qaic_config) + + def _set_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + qaic_config = self._copy_qaic_config(qaic_config) + if not hasattr(self, "hash_params"): + self.hash_params = {} + if qaic_config is not None: + pretrained_model_name_or_path = getattr(self.model, "pretrained_path", None) or self.hash_params.get( + "pretrained_model_name_or_path", None + ) + if pretrained_model_name_or_path is not None: + qaic_config.setdefault("pretrained_model_name_or_path", pretrained_model_name_or_path) + self._qaic_config = qaic_config + setattr(self.model, "qaic_config", qaic_config) + if qaic_config is None: + self.hash_params.pop("qaic_config", None) + else: + self.hash_params["qaic_config"] = qaic_config + return qaic_config + + def _resolve_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + if qaic_config is None: + qaic_config = getattr(self, "_qaic_config", None) + return self._set_qaic_config(qaic_config) def _normalize_torch_dtype(self): """ @@ -1152,6 +1185,7 @@ def transform( qaic_config: Optional[dict] = None, **compiler_options, ): + qaic_config = self._resolve_qaic_config(qaic_config) if not self.is_transformed: any_transformed = self._apply_pytorch_transforms() pooling = compiler_options.pop("pooling", getattr(self, "_pooling", None)) @@ -1160,12 +1194,12 @@ def transform( any_transformed = True any_transformed = self._post_pytorch_transform() or any_transformed + self._set_qaic_config(qaic_config) - qaic_config_for_transforms = qaic_config or getattr(self.model, "qaic_config", None) if self._supports_spd_transform(): self.model, spd_transformed = SpDTransform.apply( self.model, - qaic_config_for_transforms, + qaic_config, **compiler_options, ) self.is_tlm = getattr(self, "is_tlm", False) or spd_transformed @@ -1174,12 +1208,13 @@ def transform( if self._supports_sampler_transform(): self.model, sampler_transformed = SamplerTransform.apply( self.model, - qaic_config_for_transforms, + qaic_config, **compiler_options, ) any_transformed = any_transformed or sampler_transformed - if getattr(self, "is_tlm", False) and getattr(self.model, "qaic_config", None) is not None: - self.model.qaic_config["return_pdfs"] = True + if getattr(self, "is_tlm", False) and qaic_config is not None: + qaic_config["return_pdfs"] = True + setattr(self.model, "qaic_config", qaic_config) if not any_transformed: warnings.warn(f"No transforms applied to model: {self.model_name}. It may be an unsupported model!") @@ -1217,6 +1252,8 @@ def transform( self.hash_params.pop("blocking_kwargs", None) if qaic_config is not None: self.hash_params["qaic_config"] = qaic_config + else: + self.hash_params.pop("qaic_config", None) self.hash_params["num_replicate_kv_heads"] = effective_num_replicate_kv_heads num_cores = compiler_options.get("num_cores", compiler_options.get("aic_num_cores")) @@ -1306,6 +1343,7 @@ def _compile( """ layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) + qaic_config = self._resolve_qaic_config(qaic_config) for removed_option in ("compile_only", "compile-only"): if removed_option in compiler_options: diff --git a/QEfficient/cloud/infer.py b/QEfficient/cloud/infer.py index c1ff1248cc..2400764dd8 100644 --- a/QEfficient/cloud/infer.py +++ b/QEfficient/cloud/infer.py @@ -254,7 +254,6 @@ def main( full_batch_size=full_batch_size, local_model_dir=local_model_dir, trust_remote_code=trust_remote_code, - qaic_config=qaic_config, ) image_path = kwargs.pop("image_path", None) @@ -289,6 +288,7 @@ def main( qnn_config=qnn_config, use_onnx_subfunctions=use_onnx_subfunctions, dynamo=dynamo, + qaic_config=qaic_config, **kwargs, ) diff --git a/QEfficient/exporter/export_hf_to_cloud_ai_100.py b/QEfficient/exporter/export_hf_to_cloud_ai_100.py index 620052a4ff..d0079ac248 100644 --- a/QEfficient/exporter/export_hf_to_cloud_ai_100.py +++ b/QEfficient/exporter/export_hf_to_cloud_ai_100.py @@ -346,14 +346,14 @@ def qualcomm_efficient_converter( full_batch_size: Optional[int] = None, ) -> Tuple[str, str]: """ - This method is an alias for ``QEfficient.export``. + Deprecated public API for exporting models through the legacy ``QEfficient.export`` alias. Usage 1: This method can be used by passing ``model_name`` and ``local_model_dir`` or ``cache_dir`` if required for loading from local dir. This will download the model from ``HuggingFace`` and export it to ``ONNX`` graph and returns generated files path check below. Usage 2: You can pass ``model_name`` and ``model_kv`` as an object of ``QEfficient.QEFFAutoModelForCausalLM``, In this case will directly export the ``model_kv.model`` to ``ONNX`` - We will be deprecating this function and it will be replaced by ``QEFFAutoModelForCausalLM.export``. + This function is deprecated as a public API. Use ``QEFFAutoModelForCausalLM.from_pretrained(...).compile(...)`` instead. ``Mandatory`` Args: :model_name (str): The name of the model to be used. @@ -373,12 +373,13 @@ def qualcomm_efficient_converter( .. code-block:: python - import QEfficient - base_path, onnx_model_path = QEfficient.export(model_name="gpt2") + from QEfficient import QEFFAutoModelForCausalLM + qeff_model = QEFFAutoModelForCausalLM.from_pretrained("gpt2") + qpc_path = qeff_model.compile(num_cores=16) """ warnings.warn( - "\033[93m`qualcomm_efficient_converter` method will be deprecated soon, use `QEFFAutoModelForCausalLM.export` instead\033[0m", + "\033[93m`qualcomm_efficient_converter`/`QEfficient.export` is deprecated as a public API. Use `QEFFAutoModelForCausalLM.from_pretrained(...).compile(...)` instead.\033[0m", DeprecationWarning, stacklevel=2, ) diff --git a/QEfficient/generation/text_generation_inference.py b/QEfficient/generation/text_generation_inference.py index 17c992064c..8242cf83d1 100755 --- a/QEfficient/generation/text_generation_inference.py +++ b/QEfficient/generation/text_generation_inference.py @@ -373,11 +373,12 @@ def cloud_ai_100_exec_kv( .. code-block:: python import transformers - import QEfficient - base_path, onnx_model_path = QEfficient.export(model_name="gpt2") - qpc_path = QEfficient.compile(onnx_path=onnx_model_path, qpc_path=os.path.join(base_path, "qpc"), num_cores=14, device_group=[0]) + from QEfficient import QEFFAutoModelForCausalLM + + qeff_model = QEFFAutoModelForCausalLM.from_pretrained("gpt2") + qpc_path = qeff_model.compile(num_cores=14, num_devices=1) tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2") - exec_info = QEfficient.cloud_ai_100_exec_kv(tokenizer=tokenizer, qpc_path=qpc_path, prompt="Hi there!!", device_id=[0]) + exec_info = qeff_model.generate(tokenizer=tokenizer, prompts=["Hi there!!"], device_id=[0]) """ batch_size, ctx_len, full_batch_size = get_compilation_dims(qpc_path) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 4cfbec16e3..2afcd141be 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -8,6 +8,7 @@ import math import os import warnings +from copy import deepcopy from pathlib import Path from time import perf_counter from typing import List, Optional, Union @@ -137,6 +138,18 @@ def _resolve_torch_dtype(kwargs: dict) -> None: kwargs["dtype"] = kwargs["torch_dtype"] +def _pop_deprecated_from_pretrained_qaic_config(kwargs: dict) -> Optional[dict]: + qaic_config = kwargs.pop("qaic_config", None) + if qaic_config is not None: + warnings.warn( + "Passing `qaic_config` to `from_pretrained()` is deprecated and will be removed in a future " + "release. Pass `qaic_config` to `compile()` instead.", + DeprecationWarning, + stacklevel=3, + ) + return qaic_config + + def _ignore_public_mdp_ts_num_devices(compiler_options: dict) -> None: if "mdp_ts_num_devices" not in compiler_options: return @@ -346,6 +359,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs): QEFFTransformersBase An instance of the specific QEFFAutoModel subclass, initialized with the pretrained weights. """ + _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: @@ -484,7 +498,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k Load a QEfficient transformer model from a pretrained HuggingFace model or local path. This is the recommended way to initialize a QEfficient transformer model. The interface is similar to - ``transformers.AutoModel.from_pretrained``. Once initialized, you can use methods such as ``export``, ``compile``, and ``generate``. + ``transformers.AutoModel.from_pretrained``. Once initialized, use ``compile`` to export/compile the model and ``generate`` to run it. Parameters ---------- @@ -511,6 +525,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k QEFFAutoModel An instance initialized with the pretrained weights. """ + _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: @@ -894,6 +909,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): QEFFAutoModelForSequenceClassification An instance initialized with the pretrained weights. """ + _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: @@ -1260,26 +1276,27 @@ def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): Additional keyword arguments passed to the base class constructor. """ _configure_proxy_for_model(self, kwargs.pop("enable_proxy", False)) - self._qaic_config = qaic_config - super().__init__(model, **kwargs) + super().__init__(model, qaic_config=qaic_config, **kwargs) if hasattr(model, "get_qeff_language_decoder"): self.model = model.get_qeff_language_decoder() - self.model.qaic_config = qaic_config + self._apply_qaic_config(qaic_config) self.hash_params["qeff_auto_class"] = self.__class__.__name__ self.continuous_batching = False - if qaic_config: - if mla_absorption := qaic_config.get("mla_absorption", None): - self.hash_params["mla_absorption"] = mla_absorption - if language_model := getattr(self.model, "language_model", None): - setattr(language_model, "mla_absorption", mla_absorption) + + def _apply_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + qaic_config = self._resolve_qaic_config(qaic_config) + if qaic_config and (mla_absorption := qaic_config.get("mla_absorption", None)): + self.hash_params["mla_absorption"] = mla_absorption + if language_model := getattr(self.model, "language_model", None): + setattr(language_model, "mla_absorption", mla_absorption) + else: + self.hash_params.pop("mla_absorption", None) + return qaic_config def _post_pytorch_transform(self) -> bool: if hasattr(self.model, "get_qeff_language_decoder"): self.model = self.model.get_qeff_language_decoder() - self.model.qaic_config = self._qaic_config - if self._qaic_config and (mla_absorption := self._qaic_config.get("mla_absorption", None)): - if language_model := getattr(self.model, "language_model", None): - setattr(language_model, "mla_absorption", mla_absorption) + self._apply_qaic_config(self._qaic_config) return True return False @@ -1321,7 +1338,7 @@ def export( """ reject_legacy_moe_prefill_packed_chunk_size(kwargs) skip_transform = kwargs.pop("_qeff_skip_transform", False) - qaic_config = kwargs.pop("qaic_config", getattr(self.model, "qaic_config", None)) + qaic_config = self._apply_qaic_config(kwargs.pop("qaic_config", None)) if not skip_transform: self.transform( prefill_only=prefill_only, @@ -1450,8 +1467,6 @@ def __init__( ---------- model : nn.Module The full HuggingFace multimodal model. - qaic_config : dict, optional - A dictionary for QAIC-specific configurations. **kwargs : Additional keyword arguments. """ @@ -1467,15 +1482,28 @@ def __init__( self.vision_model = QEffVisionEncoderForTextImageToTextModel(model, **kwargs) self.lang_model = QEffCausalLMForTextImageToTextModel(model, qaic_config=qaic_config, **kwargs) self.continuous_batching = continuous_batching - self.ccl_enabled = False - if qaic_config: - self.ccl_enabled = qaic_config.get("ccl_enabled", False) + self.ccl_enabled = bool(self.lang_model._qaic_config and self.lang_model._qaic_config.get("ccl_enabled", False)) self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.input_shapes, self.output_names = None, None + def _resolve_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + if hasattr(self.lang_model, "_resolve_qaic_config"): + qaic_config = self.lang_model._resolve_qaic_config(qaic_config) + else: + if qaic_config is None: + qaic_config = getattr(self, "_qaic_config", None) + elif not isinstance(qaic_config, dict): + raise TypeError(f"`qaic_config` must be a dictionary, got {type(qaic_config).__name__}.") + else: + qaic_config = deepcopy(qaic_config) + setattr(self.lang_model, "_qaic_config", qaic_config) + self._qaic_config = qaic_config + self.ccl_enabled = bool(qaic_config and qaic_config.get("ccl_enabled", False)) + return qaic_config + @classmethod - def from_pretrained(cls, pretrained_model_name_or_path: str, qaic_config: Optional[dict] = None, **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs): """ Load a QEfficient multimodal model for dual QPC from a pretrained HuggingFace model or local path. @@ -1493,6 +1521,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, qaic_config: Option _QEffAutoModelForImageTextToTextDualQPC An instance initialized with the pretrained weights. """ + qaic_config = _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: @@ -1572,6 +1601,7 @@ def export( """ layerwise_cache_probe = kwargs.pop("_layerwise_cache_probe", False) reject_legacy_moe_prefill_packed_chunk_size(kwargs) + qaic_config = self._resolve_qaic_config(kwargs.pop("qaic_config", None)) if layerwise: return self._run_layerwise_export( export_dir=export_dir, @@ -1583,11 +1613,11 @@ def export( enable_chunking=enable_chunking, layerwise_window_size=layerwise_window_size, kv_cache_prefix=kv_cache_prefix, + qaic_config=qaic_config, **kwargs, ) bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len: int = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN - qaic_config = kwargs.get("qaic_config", getattr(self.lang_model.model, "qaic_config", None)) # TODO: move this to a DA Serving utility class if not kwargs.pop("_qeff_skip_transform", False): self.transform( @@ -1910,6 +1940,8 @@ def compile( If True, skips compilation of the language decoder. Default is False. use_onnx_subfunctions: bool, optional whether to enable ONNX subfunctions during export. Exporting PyTorch model to ONNX with modules as subfunctions helps to reduce export/compile time. Defaults to False + qaic_config : dict, optional + QAIC-specific transform configuration. Pass cache blocking, MLA, MoE, and related options here instead of to ``from_pretrained()``. **compiler_options : dict Additional compiler options for QAIC or QNN compilers. Use ``mdp_num_partitions`` to select the number of pipeline-parallel @@ -1931,6 +1963,7 @@ def compile( raise ValueError("Expected at least one of 'skip_lang' or 'skip_vision' to be False") reject_legacy_moe_prefill_packed_chunk_size(compiler_options) _ignore_public_mdp_ts_num_devices(compiler_options) + qaic_config = self._resolve_qaic_config(qaic_config) if layerwise: if skip_lang and not skip_vision: @@ -2663,12 +2696,11 @@ def __init__( "full_batch_size argument is deprecated. Use continuous_batching=True instead.", DeprecationWarning, 2 ) raise NotImplementedError("Continuous batching is not supported for image-text-to-text models yet.") - if qaic_config is not None and qaic_config.pop("include_sampler", False): + if qaic_config is not None and qaic_config.get("include_sampler", False): raise NotImplementedError("On-device sampling is not supported for single QPC multimodal models yet.") - super().__init__(model, **kwargs) - - self.model.qaic_config = qaic_config + super().__init__(model, qaic_config=qaic_config, **kwargs) + qaic_config = self._qaic_config # to handle internvl models if hasattr(self.model.config, "llm_config") and hasattr(self.model.config, "vision_config"): @@ -2681,16 +2713,13 @@ def __init__( else: self.model.config.use_cache = True self.hash_params["qeff_auto_class"] = self.__class__.__name__ - self.ccl_enabled = False - if qaic_config: - self.ccl_enabled = qaic_config.get("ccl_enabled", False) + self.ccl_enabled = bool(qaic_config and qaic_config.get("ccl_enabled", False)) self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None @classmethod def from_pretrained( cls, pretrained_model_name_or_path, - qaic_config: Optional[dict] = None, *args, **kwargs, ): @@ -2714,6 +2743,7 @@ def from_pretrained( _QEFFAutoModelForImageTextToTextSingleQPC An instance initialized with the pretrained weights. """ + qaic_config = _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: @@ -2768,13 +2798,16 @@ def export( """ reject_legacy_moe_prefill_packed_chunk_size(kwargs) skip_transform = kwargs.pop("_qeff_skip_transform", False) + qaic_config = self._resolve_qaic_config(kwargs.pop("qaic_config", None)) + if qaic_config is not None and qaic_config.get("include_sampler", False): + raise NotImplementedError("On-device sampling is not supported for single QPC multimodal models yet.") if not skip_transform: self.transform( prefill_only=prefill_only, enable_chunking=enable_chunking, prefill_seq_len=prefill_seq_len, retain_full_kv=kwargs.get("retain_full_kv", False), - qaic_config=kwargs.get("qaic_config", getattr(self.model, "qaic_config", None)), + qaic_config=qaic_config, ) inputs = self.model.get_dummy_inputs(comp_ctx_lengths=self.comp_ctx_lengths_decode) @@ -2848,6 +2881,8 @@ def compile( Not supported for this model; must be None. use_onnx_subfunctions: bool, optional whether to enable ONNX subfunctions during export. Exporting PyTorch model to ONNX with modules as subfunctions helps to reduce export/compile time. Defaults to False + qaic_config : dict, optional + QAIC-specific transform configuration. Pass cache blocking, MLA, MoE, and related options here instead of to ``from_pretrained()``. **compiler_options : dict Additional compiler options for QAIC or QNN compilers. Use ``mdp_num_partitions`` to select the number of pipeline-parallel @@ -2864,6 +2899,10 @@ def compile( If `full_batch_size`, `kv_cache_batch_size`, or `num_speculative_tokens` are not None. """ _ignore_public_mdp_ts_num_devices(compiler_options) + qaic_config = self._resolve_qaic_config(qaic_config) + self.ccl_enabled = bool(qaic_config and qaic_config.get("ccl_enabled", False)) + if qaic_config is not None and qaic_config.get("include_sampler", False): + raise NotImplementedError("On-device sampling is not supported for single QPC multimodal models yet.") if any(param is not None for param in [full_batch_size, kv_cache_batch_size, num_speculative_tokens]): raise ValueError( f"Expected 'full_batch_size', 'kv_cache_batch_size', 'num_speculative_tokens' to be None but got: " @@ -3310,7 +3349,6 @@ def from_pretrained( pretrained_model_name_or_path: str, kv_offload: Optional[bool] = None, continuous_batching: bool = False, - qaic_config: Optional[dict] = None, layerwise: bool = False, **kwargs, ): @@ -3325,8 +3363,6 @@ def from_pretrained( If True, uses the dual QPC approach (vision encoder KV offloaded). If False, uses the single QPC approach (entire model in one QPC). If None, the default behavior of the internal classes is used (typically dual QPC). - qaic_config : dict, optional - A dictionary for QAIC-specific configurations. **kwargs : Additional arguments passed to HuggingFace's ``from_pretrained``. @@ -3343,6 +3379,7 @@ def from_pretrained( NotImplementedError If `continuous_batching` is provided as True. """ + qaic_config = _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) # TODO: add a check to see if kv_offload is allowed for given model by loading the config and checking architecture or type of config here. @@ -3468,16 +3505,6 @@ def __init__( continuous_batching : bool, optional If True, enables continuous batching mode for future compilation and execution. This setting must be consistent across `from_pretrained` and `compile` calls. Default is False. - qaic_config : dict, optional - A dictionary for QAIC-specific configurations. Supported keys include: - - **speculative_model_type** (str): Specifies the type of Speculative Decoding model (e.g., "target"). - - **include_sampler** (bool): If True, enables on-device sampling of next tokens. - - **return_pdfs** (bool): If True, returns probability distributions along with sampled tokens. - For Speculative Decoding Target Language Models, this is always True. - - **max_top_k_ids** (int): Maximum number of top K tokens (<= vocab size) to consider during sampling. - - **include_guided_decoding** (bool): If True, enables guided token-level filtering - during decoding. Only works when include_sampler=True. - - **num_kv_blocks** (int): Number of K/V blocks for BlockedKV attention implementation. **kwargs : Additional keyword arguments passed to the base class constructor. @@ -3510,20 +3537,24 @@ def __init__( super().__init__(model, qaic_config=qaic_config, **kwargs) self.num_layers = model.config.num_hidden_layers self.continuous_batching = continuous_batching - self.model.qaic_config = qaic_config self.model.pretrained_path = kwargs.pop("pretrained_model_name_or_path", None) - self.is_tlm = bool(qaic_config and qaic_config.get("speculative_model_type") is not None) + self._activate_qaic_config(qaic_config) self.hash_params["qeff_auto_class"] = self.__class__.__name__ - self.ccl_enabled = False - if qaic_config: - self.ccl_enabled = qaic_config.get("ccl_enabled", False) - if mla_absorption := qaic_config.get("mla_absorption", None): - self.hash_params["mla_absorption"] = mla_absorption - setattr(self.model, "mla_absorption", mla_absorption) self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None self.hash_params["max_seq_len_cached"] = max_seq_len_cached + def _activate_qaic_config(self, qaic_config: Optional[dict]) -> Optional[dict]: + qaic_config = self._resolve_qaic_config(qaic_config) + self.is_tlm = bool(qaic_config and qaic_config.get("speculative_model_type") is not None) + self.ccl_enabled = bool(qaic_config and qaic_config.get("ccl_enabled", False)) + if qaic_config and (mla_absorption := qaic_config.get("mla_absorption", None)): + self.hash_params["mla_absorption"] = mla_absorption + setattr(self.model, "mla_absorption", mla_absorption) + else: + self.hash_params.pop("mla_absorption", None) + return qaic_config + def __repr__(self) -> str: return self.__class__.__name__ + "\n" + self.model.__repr__() @@ -3533,7 +3564,6 @@ def from_pretrained( cls, pretrained_model_name_or_path, continuous_batching: bool = False, - qaic_config: Optional[dict] = None, max_seq_len_cached: Optional[int] = None, layerwise: bool = False, weight_free: bool = False, @@ -3555,19 +3585,6 @@ def from_pretrained( Whether this model will be used for continuous batching in the future. If not set to True here, the model cannot be exported/compiled for continuous batching later. Default is False. - qaic_config : dict, optional - QAIC config dictionary. Supported keys include: - - - **speculative_model_type** (str): Specify Speculative Decoding Target Language Models. - - **include_sampler** (bool): Enable/Disable sampling of next tokens. - - **return_pdfs** (bool): Return probability distributions along with sampled next tokens. - For Speculative Decoding Target Language Model, ``return_pdfs=True`` always. - Otherwise, ``return_pdfs=True`` for Speculative Decoding Draft Language Model - and ``return_pdfs=False`` for regular model. - - **max_top_k_ids** (int): Maximum number of top K tokens (<= vocab size) to consider during sampling. - The values provided in ``top_ks`` tensor must be less than this maximum limit. - - **include_guided_decoding** (bool): If True, enables guided token-level filtering - during decoding. Only works when include_sampler=True. weight_free : bool, optional If True, builds the model on the meta device instead of loading real checkpoint weights — no weights are materialized into RAM. This is @@ -3595,6 +3612,8 @@ def from_pretrained( QEFFAutoModelForCausalLM An instance initialized with the pretrained weights. """ + qaic_config = _pop_deprecated_from_pretrained_qaic_config(kwargs) + if layerwise and weight_free: raise ValueError( "`layerwise=True` and `weight_free=True` are mutually exclusive; weight_free replaces layerwise mode." @@ -3646,9 +3665,6 @@ def from_pretrained( model = _build_meta_model(cls._hf_auto_class, pretrained_model_name_or_path, kwargs) else: model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) - if qaic_config is not None: - qaic_config["pretrained_model_name_or_path"] = pretrained_model_name_or_path - # This is support models that should be classified to in a different auto class but transformers load them via this class kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {}) if model.__class__.__name__ in MISCLASSIFIED_CAUSAL_LM_TO_QEFF_AUTO_CLASS_MAP: @@ -3801,7 +3817,7 @@ def export( self.model.config, prefill_only, kwargs.get("enable_chunking", False) ) kwargs["enable_chunking"] = enable_chunking - qaic_config = kwargs.pop("qaic_config", getattr(self.model, "qaic_config", None)) + qaic_config = self._activate_qaic_config(kwargs.pop("qaic_config", None)) # Weight-free export always uses the dynamo (torch.export) path. # Must be set here — @export_wrapper reads dynamo from kwargs before _export() body runs. dynamo = dynamo or self._weight_free @@ -3929,8 +3945,8 @@ def export( 2: "ctx_len", } output_names = [] - if self.model.qaic_config is not None and self.model.qaic_config.get("include_sampler", False): - if self.model.qaic_config.get("return_pdfs", False): + if qaic_config is not None and qaic_config.get("include_sampler", False): + if qaic_config.get("return_pdfs", False): output_names.append("probs") output_names.append("next_tokens") else: @@ -3979,8 +3995,8 @@ def export( output_names.append(f"past_{kv}.{i}_RetainedState") if "DeepseekV3ForCausalLM" in (getattr(self.model.config, "architectures", None) or []): - if self.model.qaic_config is not None and self.model.qaic_config.get("mla_absorption", None) is not None: - mla_absorption = self.model.qaic_config["mla_absorption"] + if qaic_config is not None and qaic_config.get("mla_absorption", None) is not None: + mla_absorption = qaic_config["mla_absorption"] cache_compressed = mla_absorption.get("cache_compressed", False) else: cache_compressed = False @@ -4028,14 +4044,14 @@ def export( example_inputs["num_logits_to_keep"] = torch.arange(nlk).view(nlk, 1) dynamic_axes["num_logits_to_keep"] = {0: "num_logits_to_keep"} - if self.model.qaic_config is not None and self.model.qaic_config.get("include_sampler", False): + if qaic_config is not None and qaic_config.get("include_sampler", False): example_inputs, output_names, dynamic_axes = get_sampling_inputs_and_outputs( example_inputs=example_inputs, output_names=output_names, dynamic_axes=dynamic_axes, continuous_batching=self.continuous_batching, vocab_size=self.model.config.vocab_size, - qaic_config=self.model.qaic_config, + qaic_config=qaic_config, ) # transformers>=5.3 Gemma3 models require Cache I/O internally; keep tensor/list @@ -4266,6 +4282,7 @@ def compile( mxfp6_matmul: bool = False, mxint8_kv_cache: bool = False, num_speculative_tokens: Optional[Union[int, List[int]]] = None, + qaic_config: Optional[dict] = None, prefill_only: Optional[bool] = None, use_onnx_subfunctions: bool = False, offload_pt_weights: Optional[bool] = True, @@ -4315,8 +4332,11 @@ def compile( A plain int K is treated as ``[K]`` (backward compatible). Each value K generates a decode specialization with seq_len=K+1 and num_logits_to_keep=K+1. Include 0 to compile a cheap single-token fallback - (e.g. ``[0, 3]`` for a fallback + full K=3 decode). Required if the model is - configured as a Target Language Model (``is_tlm=True``). + (e.g. ``[0, 3]`` for a fallback + full K=3 decode). Required when + ``qaic_config`` configures this model as a Target Language Model. + qaic_config : dict, optional + QAIC-specific transform configuration. Pass speculative decoding, sampler, + prefix/cache blocking, MLA, and MoE options here instead of to ``from_pretrained()``. prefill_only : bool, optional If True, compiles only for the prefill stage. If False, compiles only for the decode stage. If None, compiles for both stages. Default is None. @@ -4361,6 +4381,7 @@ def compile( """ reject_legacy_moe_prefill_packed_chunk_size(compiler_options) _ignore_public_mdp_ts_num_devices(compiler_options) + qaic_config = self._activate_qaic_config(qaic_config) enable_chunking = override_gptoss_prefill_chunking(self.model.config, prefill_only, enable_chunking) if layerwise: warnings.warn( @@ -4391,19 +4412,16 @@ def compile( offload_pt_weights=offload_pt_weights, enable_chunking=enable_chunking, retain_full_kv=retain_full_kv, + qaic_config=qaic_config, kv_cache_prefix=kv_cache_prefix, **compiler_options, ) - if self.model.qaic_config is not None and self.model.qaic_config.get("mla_absorption", None) is not None: - mla_absorption = self.model.qaic_config["mla_absorption"] + if qaic_config is not None and qaic_config.get("mla_absorption", None) is not None: + mla_absorption = qaic_config["mla_absorption"] cache_compressed = mla_absorption.get("cache_compressed", False) else: cache_compressed = False - if ( - self.model.qaic_config is not None - and self.model.qaic_config.get("mla_absorption", None) is not None - and not cache_compressed - ): + if qaic_config is not None and qaic_config.get("mla_absorption", None) is not None and not cache_compressed: logger.warning("mla_absorption will be ignored as cache_compressed is set to False") if (kv_cache_batch_size or full_batch_size) and not self.continuous_batching: logger.warning( @@ -4476,8 +4494,8 @@ def compile( _decode_ks = [validated_k] if ( - self.model.qaic_config is not None - and self.model.qaic_config.get("include_sampler", False) + qaic_config is not None + and qaic_config.get("include_sampler", False) and _decode_ks is not None and max(_decode_ks) > 0 ): @@ -4637,6 +4655,7 @@ def compile( offload_pt_weights=offload_pt_weights, enable_chunking=enable_chunking, retain_full_kv=retain_full_kv, + qaic_config=qaic_config, kv_cache_prefix=kv_cache_prefix, **compiler_options, ) @@ -5203,6 +5222,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k # You can now execute the model out = model.generate(processor,inputs=input_audio) """ + _pop_deprecated_from_pretrained_qaic_config(kwargs) enable_proxy = kwargs.pop("enable_proxy", False) if kwargs.get("attn_implementation", None) not in {None, "eager"}: logger.warning('Updating attn_implementation="eager"') diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index 78bf2675be..c3329f8bc6 100644 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -283,7 +283,7 @@ def export_wrapper(func): def wrapper(self, *args, **kwargs): if not _EXPORT_FROM_COMPILE.get(): warnings.warn( - "Direct .export() is deprecated. Use .compile() to export and compile with the complete configuration.", + "Direct .export() is deprecated as a public API. Use .compile() to export and compile with the complete configuration.", DeprecationWarning, stacklevel=2, ) diff --git a/QEfficient/utils/test_utils.py b/QEfficient/utils/test_utils.py index 52a692b965..ea40eea3e6 100644 --- a/QEfficient/utils/test_utils.py +++ b/QEfficient/utils/test_utils.py @@ -76,7 +76,7 @@ def load_qeff_causal_lm_model( qaic_config: Dict = None, config: Optional[AutoConfig] = None, ): - kwargs = dict(continuous_batching=continuous_batching, qaic_config=qaic_config) + kwargs = dict(continuous_batching=continuous_batching) if config is None: if num_hidden_layers != -1: kwargs["num_hidden_layers"] = num_hidden_layers @@ -84,6 +84,8 @@ def load_qeff_causal_lm_model( else: model_hf = load_hf_causal_lm_model(model_name, num_hidden_layers, config) qeff_model = QEFFAutoModelForCausalLM(model_hf, **kwargs) + if qaic_config is not None: + qeff_model._set_qaic_config(qaic_config) return qeff_model @@ -185,7 +187,6 @@ def load_qeff_vlm_model( config=config, kv_offload=kv_offload, continuous_batching=continuous_batching, - qaic_config=qaic_config, torch_dtype=torch_dtype, ignore_mismatched_sizes=True, trust_remote_code=True, @@ -197,7 +198,6 @@ def load_qeff_vlm_model( config=config, kv_offload=kv_offload, continuous_batching=continuous_batching, - qaic_config=qaic_config, trust_remote_code=True, ignore_mismatched_sizes=True, torch_dtype=torch_dtype, @@ -212,10 +212,11 @@ def load_qeff_vlm_model( kv_offload=kv_offload, continuous_batching=continuous_batching, trust_remote_code=True, - qaic_config=qaic_config, torch_dtype=torch_dtype, ) + if qaic_config is not None: + qeff_model._set_qaic_config(qaic_config) return qeff_model diff --git a/docs/source/features_enablement.md b/docs/source/features_enablement.md index f769c102de..a6c74b8a25 100644 --- a/docs/source/features_enablement.md +++ b/docs/source/features_enablement.md @@ -108,7 +108,7 @@ generated_qpc_path = qeff_model.compile( ## Draft-Based Speculative Decoding Draft-based speculative decoding is a technique where a small Draft Language Model (DLM) makes `num_speculative_tokens` autoregressive speculations ahead of the Target Language Model (TLM). The objective is to predict what the TLM would have predicted if it would have been used instead of the DLM. This approach is beneficial when the autoregressive decode phase of the TLM is memory bound and thus, we can leverage the extra computing resources of our hardware by batching the speculations of the DLM as an input to TLM to validate the speculations. -To export and compile both DLM/TLM, add corresponding `qaic_config` and `num_speculative_tokens` for TLM and export DLM as you would any other QEfficient LLM model: +To export and compile both DLM/TLM, pass the TLM `qaic_config` and `num_speculative_tokens` to `compile()` and compile the DLM as you would any other QEfficient LLM model: ```Python from QEfficient import QEFFAutoModelForCausalLM as AutoModelForCausalLM @@ -117,10 +117,10 @@ tlm_name = "meta-llama/Llama-2-70b-chat-hf" dlm_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" k = 3 # DLM will make `k` speculations qaic_config = dict(speculative_model_type="target") -tlm = AutoModelForCausalLM.from_pretrained(tlm_name, qaic_config=qaic_config) +tlm = AutoModelForCausalLM.from_pretrained(tlm_name) dlm = AutoModelForCausalLM.from_pretrained(dlm_name) -tlm.compile(num_speculative_tokens=k) +tlm.compile(qaic_config=qaic_config, num_speculative_tokens=k) dlm.compile() ``` -The `qaic_config` dictionary is fed during the instantiation of the model because slight changes to the ONNX graph are required. Once complete, the user can specify `num_speculative_tokens` to define the actual number of speculations that the TLM will take as input during the decode phase. As for the DLM, no new changes are required at the ONNX or compile level. +Pass the `qaic_config` dictionary to `compile()` because compile drives the ONNX export and applies the graph changes required for speculative decoding. `num_speculative_tokens` defines how many speculations the TLM accepts during decode. As for the DLM, no new changes are required at the ONNX or compile level. diff --git a/docs/source/qeff_autoclasses.md b/docs/source/qeff_autoclasses.md index 0d065146cb..628e044d70 100644 --- a/docs/source/qeff_autoclasses.md +++ b/docs/source/qeff_autoclasses.md @@ -14,7 +14,6 @@ ```{eval-rst} .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCausalLM.from_pretrained -.. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCausalLM.export .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCausalLM.compile .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCausalLM.generate ``` @@ -40,7 +39,6 @@ Do not pass `mdp_ts_num_devices` to this public `compile()` API. It is ignored w ```{eval-rst} .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModel.from_pretrained -.. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModel.export .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModel.compile .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModel.generate ``` @@ -60,7 +58,6 @@ Do not pass `mdp_ts_num_devices` to this public `compile()` API. It is ignored w ```{eval-rst} .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSequenceClassification.from_pretrained -.. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSequenceClassification.export .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSequenceClassification.compile .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSequenceClassification.generate ``` @@ -80,7 +77,6 @@ Do not pass `mdp_ts_num_devices` to this public `compile()` API. It is ignored w ```{eval-rst} .. automethod:: QEfficient.peft.auto.QEffAutoPeftModelForCausalLM.from_pretrained -.. automethod:: QEfficient.peft.auto.QEffAutoPeftModelForCausalLM.export .. automethod:: QEfficient.peft.auto.QEffAutoPeftModelForCausalLM.compile .. automethod:: QEfficient.peft.auto.QEffAutoPeftModelForCausalLM.generate ``` @@ -100,7 +96,6 @@ Do not pass `mdp_ts_num_devices` to this public `compile()` API. It is ignored w ```{eval-rst} .. automethod:: QEfficient.peft.lora.auto.QEffAutoLoraModelForCausalLM.from_pretrained -.. automethod:: QEfficient.peft.lora.auto.QEffAutoLoraModelForCausalLM.export .. automethod:: QEfficient.peft.lora.auto.QEffAutoLoraModelForCausalLM.compile .. automethod:: QEfficient.peft.lora.auto.QEffAutoLoraModelForCausalLM.generate ``` @@ -143,7 +138,6 @@ Do not pass `mdp_ts_num_devices` to the public auto-model `compile()` API. It is ```{eval-rst} .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSpeechSeq2Seq.from_pretrained -.. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSpeechSeq2Seq.export .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSpeechSeq2Seq.compile .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForSpeechSeq2Seq.generate ``` @@ -163,7 +157,6 @@ Do not pass `mdp_ts_num_devices` to the public auto-model `compile()` API. It is ```{eval-rst} .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCTC.from_pretrained -.. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCTC.export .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCTC.compile .. automethod:: QEfficient.transformers.models.modeling_auto.QEFFAutoModelForCTC.generate ``` diff --git a/docs/source/quick_start.md b/docs/source/quick_start.md index 35a270a491..318d256334 100644 --- a/docs/source/quick_start.md +++ b/docs/source/quick_start.md @@ -30,12 +30,12 @@ Use ``bash terminal``, else if using ``ZSH terminal`` then ``device_group``shoul Below are the Command Line APIs we support for infernce in the library. #### Export -**CLI API:** [`QEfficient.cloud.export`](#export_api) +**Deprecated CLI API:** [`QEfficient.cloud.export`](#export_api) -User can export a model to ONNX using the CLI command. This will convert the model to an ONNX format and store the resulting ONNX model file in the QEfficient cache folder. [Click here](#export_api) for more information about the export command and arguments explanation. +Direct export is deprecated as an external user-facing API. Use `QEfficient.cloud.infer` or the Python `.compile()` API instead; both paths export ONNX as needed with the complete compile configuration. ```bash -python -m QEfficient.cloud.export --model_name gpt2 +python -m QEfficient.cloud.infer --model_name gpt2 --batch_size 1 --prompt_len 32 --ctx_len 128 --num_cores 16 --device_group [0] --prompt "My name is" ``` --- @@ -197,7 +197,7 @@ print(f"{model_name} optimized for Cloud AIxxx (AI100, AI200 and so on) \n", qef ### 2. Export and Compile with one API -Use the qualcomm_efficient_converter API to export the KV transformed Model to ONNX and Verify on Torch. +Use the model `compile()` API to export the transformed model to ONNX and compile it for Cloud AI hardware. The legacy `qualcomm_efficient_converter`/`QEfficient.export` entry points are deprecated as public APIs. ```Python # We can now export the modified models to ONNX framework diff --git a/examples/embeddings/qwen3vl/qwen3_vl_embedding.py b/examples/embeddings/qwen3vl/qwen3_vl_embedding.py index bd707ffb08..34c0e9b4c0 100644 --- a/examples/embeddings/qwen3vl/qwen3_vl_embedding.py +++ b/examples/embeddings/qwen3vl/qwen3_vl_embedding.py @@ -110,7 +110,6 @@ def main() -> None: kv_offload=True, trust_remote_code=True, config=config, - qaic_config={"export_embedding": True}, ) # 2) Build embedding helper and reference payload. @@ -135,6 +134,7 @@ def main() -> None: num_cores=args.num_cores, num_devices=args.num_devices, mxfp6_matmul=args.mxfp6_matmul, + qaic_config={"export_embedding": True}, ) # 5) Run AI100 embedding generation on precompiled QPCs. diff --git a/examples/image_text_to_text/models/kimi_k2/example_kimi_k25_vision_disagg.py b/examples/image_text_to_text/models/kimi_k2/example_kimi_k25_vision_disagg.py index 579bc72629..54637734a9 100644 --- a/examples/image_text_to_text/models/kimi_k2/example_kimi_k25_vision_disagg.py +++ b/examples/image_text_to_text/models/kimi_k2/example_kimi_k25_vision_disagg.py @@ -352,7 +352,6 @@ def main(): kv_offload=True, config=model.config, torch_dtype=torch.float32, - qaic_config=qaic_config, layerwise=False, ) diff --git a/examples/image_text_to_text/models/kimi_k2/export_kimik2.py b/examples/image_text_to_text/models/kimi_k2/export_kimik2.py index 3add5b2070..28107979f8 100644 --- a/examples/image_text_to_text/models/kimi_k2/export_kimik2.py +++ b/examples/image_text_to_text/models/kimi_k2/export_kimik2.py @@ -1,53 +1,53 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer - -from QEfficient import QEFFAutoModelForCausalLM - -# parameters to be configured -prompt = "Once upon a time," -num_hidden_layers = 2 -TS = 4 -mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} -# qaic_config = None # Full PKV Cache -# qaic_config = {"blocking_mode": "h"} # Full PKV Cache with Head Blocking -# qaic_config = {"mla_absorption": mla_absorption} # for No Blocking -# qaic_config = {"mla_absorption": mla_absorption, "replicate_kv_heads": True} # No blocking with kv head replication -# qaic_config = {"mla_absorption": mla_absorption, "blocking_mode": "kv"} # for KV blocking -# qaic_config = {"mla_absorption": mla_absorption, "blocking_mode": "kv", "replicate_kv_heads": True} # for KV blocking with kv head replication -qaic_config = { - "mla_absorption": mla_absorption, - "blocking_mode": "h", - "replicate_kv_heads": True, -} -# for h blocking, it internally sets head_block_size equal to num_devices/num_replicate_kv_heads (computed) - -model_name = "moonshotai/Kimi-K2-Thinking" -model = AutoModelForCausalLM.from_pretrained( - model_name, torch_dtype=torch.float32, num_hidden_layers=num_hidden_layers, trust_remote_code=True -) -tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) - -qeff_model = QEFFAutoModelForCausalLM(model, qaic_config=qaic_config) - -prefill_seq_len = 1 -ctx_len = 16 * 1024 - -qpc_path = qeff_model.compile( - prefill_seq_len=prefill_seq_len, - ctx_len=ctx_len, - mxfp6_matmul=True, - mxint8_kv_cache=False, - num_devices=TS, - num_cores=16, - use_onnx_subfunctions=True, - qaic_config=qaic_config, -) - -qeff_model.generate(prompts=["Once upon a time,"], tokenizer=tokenizer) +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from QEfficient import QEFFAutoModelForCausalLM + +# parameters to be configured +prompt = "Once upon a time," +num_hidden_layers = 2 +TS = 4 +mla_absorption = {"cache_compressed": True, "absorption": False, "online": False} +# qaic_config = None # Full PKV Cache +# qaic_config = {"blocking_mode": "h"} # Full PKV Cache with Head Blocking +# qaic_config = {"mla_absorption": mla_absorption} # for No Blocking +# qaic_config = {"mla_absorption": mla_absorption, "replicate_kv_heads": True} # No blocking with kv head replication +# qaic_config = {"mla_absorption": mla_absorption, "blocking_mode": "kv"} # for KV blocking +# qaic_config = {"mla_absorption": mla_absorption, "blocking_mode": "kv", "replicate_kv_heads": True} # for KV blocking with kv head replication +qaic_config = { + "mla_absorption": mla_absorption, + "blocking_mode": "h", + "replicate_kv_heads": True, +} +# for h blocking, it internally sets head_block_size equal to num_devices/num_replicate_kv_heads (computed) + +model_name = "moonshotai/Kimi-K2-Thinking" +model = AutoModelForCausalLM.from_pretrained( + model_name, torch_dtype=torch.float32, num_hidden_layers=num_hidden_layers, trust_remote_code=True +) +tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + +qeff_model = QEFFAutoModelForCausalLM(model) + +prefill_seq_len = 1 +ctx_len = 16 * 1024 + +qpc_path = qeff_model.compile( + prefill_seq_len=prefill_seq_len, + ctx_len=ctx_len, + mxfp6_matmul=True, + mxint8_kv_cache=False, + num_devices=TS, + num_cores=16, + use_onnx_subfunctions=True, + qaic_config=qaic_config, +) + +qeff_model.generate(prompts=["Once upon a time,"], tokenizer=tokenizer) diff --git a/examples/image_text_to_text/models/qwen3vl/qwen3_vl_blocked.py b/examples/image_text_to_text/models/qwen3vl/qwen3_vl_blocked.py index 1308acdef0..e94f91b363 100644 --- a/examples/image_text_to_text/models/qwen3vl/qwen3_vl_blocked.py +++ b/examples/image_text_to_text/models/qwen3vl/qwen3_vl_blocked.py @@ -35,7 +35,6 @@ attn_implementation="eager", kv_offload=True, config=config, - qaic_config=qaic_config, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) @@ -60,6 +59,7 @@ skip_vision=True, mos=1, use_onnx_subfunctions=False, + qaic_config=qaic_config, ) messages = [ @@ -105,6 +105,7 @@ aic_enable_depth_first=True, mos=1, use_onnx_subfunctions=False, + qaic_config=qaic_config, ) ### IMAGE + TEXT ### diff --git a/examples/performance/compute_context_length/basic_inference.py b/examples/performance/compute_context_length/basic_inference.py index 6e8c045fbc..8c3259ab48 100644 --- a/examples/performance/compute_context_length/basic_inference.py +++ b/examples/performance/compute_context_length/basic_inference.py @@ -109,13 +109,10 @@ def main(): print(f"Loading model: {args.model_name}") print(f"Continuous batching: {args.continuous_batching}") - # Load model with CCL configuration + # Load model model = QEFFAutoModelForCausalLM.from_pretrained( args.model_name, continuous_batching=args.continuous_batching, - qaic_config={ - "ccl_enabled": args.ccl_enabled, - }, ) # Compile the model @@ -134,6 +131,7 @@ def main(): if args.ccl_enabled: compile_kwargs["comp_ctx_lengths_prefill"] = args.comp_ctx_lengths_prefill compile_kwargs["comp_ctx_lengths_decode"] = args.comp_ctx_lengths_decode + compile_kwargs["qaic_config"] = {"ccl_enabled": args.ccl_enabled} qpc_path = model.compile(**compile_kwargs) print(f"Model compiled successfully to: {qpc_path}") diff --git a/examples/performance/compute_context_length/gemma3.py b/examples/performance/compute_context_length/gemma3.py index 1dcec5c811..85f563bf76 100644 --- a/examples/performance/compute_context_length/gemma3.py +++ b/examples/performance/compute_context_length/gemma3.py @@ -20,7 +20,7 @@ tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) processor = AutoProcessor.from_pretrained(model_id) -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -30,6 +30,9 @@ ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. comp_ctx_lengths_prefill = [3072] comp_ctx_lengths_decode = [4096, ctx_len] @@ -41,9 +44,6 @@ config=config, attn_implementation="eager", kv_offload=True, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) ### use skip_vision=True, if want to run only text, or false ### @@ -65,6 +65,7 @@ node_precision_info="examples/performance/compute_context_length/fp32_nodes_gemma3_4b.yaml", comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) messages = [ @@ -103,6 +104,7 @@ node_precision_info="examples/performance/compute_context_length/fp32_nodes_gemma3_4b.yaml", comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ### IMAGE + TEXT ### diff --git a/examples/performance/compute_context_length/gpt_oss.py b/examples/performance/compute_context_length/gpt_oss.py index 92bef9148b..aaaaab7d7f 100644 --- a/examples/performance/compute_context_length/gpt_oss.py +++ b/examples/performance/compute_context_length/gpt_oss.py @@ -11,7 +11,7 @@ model_id = "openai/gpt-oss-20b" # weights are not required to convert to fp32 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -21,15 +21,15 @@ ctx_len = 4096 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. # In moe models like gpt-oss, since prefill_seq_len=1 both comp_ctx_lengths_prefill and comp_ctx_lengths_decode can share similar lists. comp_ctx_lengths_prefill = comp_ctx_lengths_decode = [1024, ctx_len] qeff_model = QEFFAutoModelForCausalLM.from_pretrained( model_id, - qaic_config={ - "ccl_enabled": True, - }, ) tokenizer = AutoTokenizer.from_pretrained(model_id) @@ -45,6 +45,7 @@ num_speculative_tokens=None, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) print(f"qpc path is {qpc_path}") streamer = TextStreamer(tokenizer) diff --git a/examples/performance/compute_context_length/gpt_oss_disagg_mode_with_chunking.py b/examples/performance/compute_context_length/gpt_oss_disagg_mode_with_chunking.py index 50f5136700..82f6f1bdf8 100644 --- a/examples/performance/compute_context_length/gpt_oss_disagg_mode_with_chunking.py +++ b/examples/performance/compute_context_length/gpt_oss_disagg_mode_with_chunking.py @@ -34,12 +34,11 @@ PREFILL_SEQ_LEN = 128 CTX_LEN = 4096 -qeff_model = QEFFAutoModelForCausalLM.from_pretrained( - model_id, - qaic_config={ - "ccl_enabled": True, - }, -) +qaic_config = { + "ccl_enabled": True, +} + +qeff_model = QEFFAutoModelForCausalLM.from_pretrained(model_id) comp_ctx_lengths_decode = [1024, 2048, 4096] @@ -57,6 +56,7 @@ retain_full_kv=True, prefill_only=False, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, # # split_retained_state_io=True, # This should be used for disagg serving via VLLM # node_precision_info=non_subfunc_npi_file_path, ) diff --git a/examples/performance/compute_context_length/granite_vision.py b/examples/performance/compute_context_length/granite_vision.py index ef5dc3a517..75dad2434d 100644 --- a/examples/performance/compute_context_length/granite_vision.py +++ b/examples/performance/compute_context_length/granite_vision.py @@ -37,13 +37,14 @@ def run_model( # The Dual QPC approach splits the model to perform Image Encoding and Output generation in 2 different QPCs. # The outputs of the Vision Encoder are then passed to the Language model via host in this case. + qaic_config = { + "ccl_enabled": ccl_enabled, + } + model = QEFFAutoModelForImageTextToText.from_pretrained( model_name, token=token, kv_offload=kv_offload, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) ## STEP - 2 Export & Compile the Model @@ -57,6 +58,7 @@ def run_model( mxfp6_matmul=False, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ## STEP - 3 Load and process the inputs for Inference diff --git a/examples/performance/compute_context_length/internvl.py b/examples/performance/compute_context_length/internvl.py index 02e965e0de..036f942593 100644 --- a/examples/performance/compute_context_length/internvl.py +++ b/examples/performance/compute_context_length/internvl.py @@ -184,13 +184,14 @@ def run_intern_on_aic( # The original Intern-VL model, despite being multimodal, is loaded using `AutoModelForCausalLM` in Huggingface. # To maintain compatibility, we load this model using `QEFFAutoModelForCausalLM`. + qaic_config = { + "ccl_enabled": ccl_enabled, + } + model = QEFFAutoModelForCausalLM.from_pretrained( model_name, kv_offload=kv_offload, trust_remote_code=True, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) ## STEP 2 -- EXPORT & COMPILE THE MODEL @@ -203,6 +204,7 @@ def run_intern_on_aic( mxfp6_matmul=False, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ## STEP 3 -- SETUP THE PROCESSOR diff --git a/examples/performance/compute_context_length/llama4.py b/examples/performance/compute_context_length/llama4.py index a867e1bd33..9457dae174 100644 --- a/examples/performance/compute_context_length/llama4.py +++ b/examples/performance/compute_context_length/llama4.py @@ -17,7 +17,7 @@ config.text_config.num_hidden_layers = 4 config.vision_config.num_hidden_layers = 2 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -27,6 +27,9 @@ ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. # Set the list of ccl during prefilling process comp_ctx_lengths_prefill = [3072] @@ -38,9 +41,6 @@ attn_implementation="eager", kv_offload=True, config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) @@ -64,6 +64,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) messages = [ @@ -107,6 +108,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ### IMAGE + TEXT ### diff --git a/examples/performance/compute_context_length/llama4_cb.py b/examples/performance/compute_context_length/llama4_cb.py index f971606931..881d5f61cc 100644 --- a/examples/performance/compute_context_length/llama4_cb.py +++ b/examples/performance/compute_context_length/llama4_cb.py @@ -1,118 +1,117 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -import transformers -from transformers import AutoConfig, AutoProcessor - -from QEfficient import QEFFAutoModelForImageTextToText - -model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" -config = AutoConfig.from_pretrained(model_id) -# For Testing Purpose Only -config.text_config.num_hidden_layers = 4 -config.vision_config.num_hidden_layers = 2 - -tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) -processor = AutoProcessor.from_pretrained(model_id) - -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). -## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. -## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. -## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. -## - The second list, comp_ctx_lengths_decode, defines the compute-context-length values for the decoding process. -## -- During decoding, the model selects an appropriate context length from the list based on the input prompt length and cache index. -## -- It starts from the correct value in the list and increases the context length dynamically when the generated token's cache index exceeds the current CCL value. - -ctx_len = 4096 -ccl_enabled = True -# Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. -# Set the list of ccl during prefilling process -comp_ctx_lengths_prefill = [3072] -# Set the list of ccl during decoding process -comp_ctx_lengths_decode = [ctx_len] - -continious_batching = True -if continious_batching: - qeff_model = QEFFAutoModelForImageTextToText.from_pretrained( - model_id, - attn_implementation="eager", - kv_offload=True, - config=config, - continuous_batching=True, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, - ) - - qeff_model.compile( - prefill_seq_len=128, - ctx_len=ctx_len, - img_size=336, - num_cores=16, - num_devices=4, - max_num_tiles=17, - batch_size=1, - full_batch_size=4, - mxfp6_matmul=True, - mxint8_kv_cache=True, - aic_enable_depth_first=True, - mos=1, - comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, - comp_ctx_lengths_decode=comp_ctx_lengths_decode, - ) -else: - qeff_model = QEFFAutoModelForImageTextToText.from_pretrained( - model_id, - attn_implementation="eager", - kv_offload=True, - config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, - ) - - qeff_model.compile( - prefill_seq_len=128, - ctx_len=ctx_len, - img_size=336, - num_cores=16, - num_devices=4, - max_num_tiles=17, - batch_size=1, - mxfp6_matmul=True, - mxint8_kv_cache=True, - aic_enable_depth_first=True, - mos=1, - comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, - comp_ctx_lengths_decode=comp_ctx_lengths_decode, - ) - -image_urls = [ - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png", - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg", - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png", - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg", -] - -prompts = [ - "Can you describe the image in detail?", - "What are the objects in the image?", - "What is the main subject of the image?", - "What colors are predominant in the image?", -] - -exec_info = qeff_model.generate( - tokenizer=tokenizer, - prompts=prompts, - processor=processor, - images=image_urls, - generation_len=100, -) - -# print("Generated texts:", exec_info.generated_texts) -print("Generated IDs:", exec_info.generated_ids) -print(exec_info) +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import transformers +from transformers import AutoConfig, AutoProcessor + +from QEfficient import QEFFAutoModelForImageTextToText + +model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" +config = AutoConfig.from_pretrained(model_id) +# For Testing Purpose Only +config.text_config.num_hidden_layers = 4 +config.vision_config.num_hidden_layers = 2 + +tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) +processor = AutoProcessor.from_pretrained(model_id) + +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). +## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. +## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. +## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. +## - The second list, comp_ctx_lengths_decode, defines the compute-context-length values for the decoding process. +## -- During decoding, the model selects an appropriate context length from the list based on the input prompt length and cache index. +## -- It starts from the correct value in the list and increases the context length dynamically when the generated token's cache index exceeds the current CCL value. + +ctx_len = 4096 +ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} +# Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. +# Set the list of ccl during prefilling process +comp_ctx_lengths_prefill = [3072] +# Set the list of ccl during decoding process +comp_ctx_lengths_decode = [ctx_len] + +continious_batching = True +if continious_batching: + qeff_model = QEFFAutoModelForImageTextToText.from_pretrained( + model_id, + attn_implementation="eager", + kv_offload=True, + config=config, + continuous_batching=True, + ) + + qeff_model.compile( + prefill_seq_len=128, + ctx_len=ctx_len, + img_size=336, + num_cores=16, + num_devices=4, + max_num_tiles=17, + batch_size=1, + full_batch_size=4, + mxfp6_matmul=True, + mxint8_kv_cache=True, + aic_enable_depth_first=True, + mos=1, + comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, + comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, + ) +else: + qeff_model = QEFFAutoModelForImageTextToText.from_pretrained( + model_id, + attn_implementation="eager", + kv_offload=True, + config=config, + ) + + qeff_model.compile( + prefill_seq_len=128, + ctx_len=ctx_len, + img_size=336, + num_cores=16, + num_devices=4, + max_num_tiles=17, + batch_size=1, + mxfp6_matmul=True, + mxint8_kv_cache=True, + aic_enable_depth_first=True, + mos=1, + comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, + comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, + ) + +image_urls = [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg", +] + +prompts = [ + "Can you describe the image in detail?", + "What are the objects in the image?", + "What is the main subject of the image?", + "What colors are predominant in the image?", +] + +exec_info = qeff_model.generate( + tokenizer=tokenizer, + prompts=prompts, + processor=processor, + images=image_urls, + generation_len=100, +) + +# print("Generated texts:", exec_info.generated_texts) +print("Generated IDs:", exec_info.generated_ids) +print(exec_info) diff --git a/examples/performance/compute_context_length/llama4_multi_image.py b/examples/performance/compute_context_length/llama4_multi_image.py index 314aa49b3e..2a9cf589bd 100644 --- a/examples/performance/compute_context_length/llama4_multi_image.py +++ b/examples/performance/compute_context_length/llama4_multi_image.py @@ -17,7 +17,7 @@ config.text_config.num_hidden_layers = 4 config.vision_config.num_hidden_layers = 2 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -27,6 +27,9 @@ ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. # Set the list of ccl during prefilling process comp_ctx_lengths_prefill = [5376] @@ -38,9 +41,6 @@ attn_implementation="eager", kv_offload=True, config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) @@ -59,6 +59,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ### Multi_image Prompt ### diff --git a/examples/performance/compute_context_length/mistral3.py b/examples/performance/compute_context_length/mistral3.py index a773ddfd94..fbe088e47c 100644 --- a/examples/performance/compute_context_length/mistral3.py +++ b/examples/performance/compute_context_length/mistral3.py @@ -42,13 +42,14 @@ def run_model( config.text_config.num_hidden_layers = 4 config.vision_config.num_hidden_layers = 2 + qaic_config = { + "ccl_enabled": ccl_enabled, + } + model = QEFFAutoModelForImageTextToText.from_pretrained( model_name, kv_offload=kv_offload, config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) ## STEP - 2 Export & Compile the Model @@ -62,6 +63,7 @@ def run_model( mxfp6_matmul=False, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ## STEP - 3 Load and process the inputs for Inference diff --git a/examples/performance/compute_context_length/molmo.py b/examples/performance/compute_context_length/molmo.py index 8d773f5fe0..dc7324b28f 100644 --- a/examples/performance/compute_context_length/molmo.py +++ b/examples/performance/compute_context_length/molmo.py @@ -18,7 +18,7 @@ # For Testing Purpose Only # config.num_hidden_layers = 2 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -29,6 +29,9 @@ # load the model ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. comp_ctx_lengths_prefill = [3072] # None # comp_ctx_lengths_decode = [4096, 8192] # None # @@ -38,9 +41,6 @@ kv_offload=True, trust_remote_code=True, config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) @@ -61,6 +61,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) inputs = processor.process(text="Tell me about yourself") @@ -86,6 +87,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ### IMAGE + TEXT ### diff --git a/examples/performance/compute_context_length/qwen2_5_vl.py b/examples/performance/compute_context_length/qwen2_5_vl.py index 5a68189306..1223d5aff7 100644 --- a/examples/performance/compute_context_length/qwen2_5_vl.py +++ b/examples/performance/compute_context_length/qwen2_5_vl.py @@ -22,7 +22,7 @@ # For Testing Purpose Only config.text_config.num_hidden_layers = 2 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -32,6 +32,9 @@ ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. comp_ctx_lengths_prefill = [4096] # None # comp_ctx_lengths_decode = [6144, ctx_len] # None # @@ -41,9 +44,6 @@ attn_implementation="eager", kv_offload=True, config=config, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) @@ -70,6 +70,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) messages = [ @@ -116,6 +117,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) ### IMAGE + TEXT ### diff --git a/examples/performance/compute_context_length/qwen2_5_vl_cb.py b/examples/performance/compute_context_length/qwen2_5_vl_cb.py index c247a1e587..91c3b94b6f 100644 --- a/examples/performance/compute_context_length/qwen2_5_vl_cb.py +++ b/examples/performance/compute_context_length/qwen2_5_vl_cb.py @@ -19,7 +19,7 @@ # For Testing Purpose Only config.text_config.num_hidden_layers = 4 -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -29,6 +29,9 @@ ctx_len = 8192 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. comp_ctx_lengths_prefill = [4096] comp_ctx_lengths_decode = [6144, ctx_len] @@ -39,9 +42,6 @@ kv_offload=True, config=config, continuous_batching=True, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) @@ -63,6 +63,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) image_urls = [ diff --git a/examples/performance/compute_context_length/qwen3moe.py b/examples/performance/compute_context_length/qwen3moe.py index 93849fa5a3..821cb80b14 100644 --- a/examples/performance/compute_context_length/qwen3moe.py +++ b/examples/performance/compute_context_length/qwen3moe.py @@ -16,7 +16,7 @@ # We will use prompt_len=1 for compilation for both cb and non-cb inference """ -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths_prefill and comp_ctx_lengths_decode to provide two lists of context lengths for the prefilling and decoding processes. If both are None, the lists will be generated automatically based on the context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -27,6 +27,9 @@ ctx_len = 1024 prefill_seq_len = 1 ccl_enabled = True +qaic_config = { + "ccl_enabled": ccl_enabled, +} # Two optional lists, comp_ctx_lengths_prefill and comp_ctx_lengths_decode, define CCL values for prefilling and decoding. # In moe models when compiling with prefill_seq_len=1 and non-continuous-batching mode, prefill and decode will share the same ccl specializations. comp_ctx_lengths_prefill = comp_ctx_lengths_decode = [256, 512, ctx_len] @@ -34,9 +37,6 @@ model = QEFFAutoModelForCausalLM.from_pretrained( model_name, continuous_batching=False, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) model.compile( @@ -50,6 +50,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) tokenizer = AutoTokenizer.from_pretrained(model_name) exec_info = model.generate(prompts=Constants.INPUT_STR, tokenizer=tokenizer) diff --git a/examples/performance/compute_context_length/qwen3moe_example/ccl_qwen3moe_inference.py b/examples/performance/compute_context_length/qwen3moe_example/ccl_qwen3moe_inference.py index 9fb4c4d43c..2f3aef7513 100644 --- a/examples/performance/compute_context_length/qwen3moe_example/ccl_qwen3moe_inference.py +++ b/examples/performance/compute_context_length/qwen3moe_example/ccl_qwen3moe_inference.py @@ -16,7 +16,7 @@ # We will use prompt_len=1 for compilation for both cb and non-cb inference """ -## Activate Compute-Context-Length (CCL) feature by setting ccl_enabled=True when loading the model with from_pretrained(). +## Activate Compute-Context-Length (CCL) feature by passing ccl_enabled=True to compile(). ## Use the optional comp_ctx_lengths argument to provide two lists of context lengths for the prefilling and decoding processes. If comp_ctx_lengths=None, the model will run with its default context length. ## - The first list, comp_ctx_lengths_prefill, defines the compute-context-length values for the prefilling process. ## -- The process starts with the first value in the list and gradually increases the context length based on the position_id of the current prompt chunk. @@ -26,6 +26,9 @@ ctx_len = 1024 prefill_seq_len = 1 +qaic_config = { + "ccl_enabled": True, +} # In moe models when compiling with prefill_seq_len=1 and non-continuous-batching mode, prefill and decode will share the same ccl specializations. comp_ctx_lengths_prefill = [256, 512, ctx_len] # None # comp_ctx_lengths_decode = [256, 512, ctx_len] # None # @@ -33,7 +36,6 @@ model = QEFFAutoModelForCausalLM.from_pretrained( model_name, continuous_batching=False, - ccl_enabled=True, num_hidden_layers=4, ) @@ -48,6 +50,7 @@ mos=1, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) # mos=1, tokenizer = AutoTokenizer.from_pretrained(model_name) diff --git a/examples/performance/compute_context_length/vlm_inference.py b/examples/performance/compute_context_length/vlm_inference.py index 294632fe3e..6ca99b194b 100644 --- a/examples/performance/compute_context_length/vlm_inference.py +++ b/examples/performance/compute_context_length/vlm_inference.py @@ -68,14 +68,15 @@ def run_model( # - Dual QPC (kv_offload=True): Vision encoder and language model run in separate QPCs # with outputs transferred via host for flexibility + qaic_config = { + "ccl_enabled": ccl_enabled, + } + model = QEFFAutoModelForImageTextToText.from_pretrained( model_name, token=hf_token, attn_implementation="eager", kv_offload=kv_offload, - qaic_config={ - "ccl_enabled": ccl_enabled, - }, ) ## STEP 2: Export & Compile the Model @@ -90,6 +91,7 @@ def run_model( mxfp6_matmul=False, comp_ctx_lengths_prefill=comp_ctx_lengths_prefill, comp_ctx_lengths_decode=comp_ctx_lengths_decode, + qaic_config=qaic_config, ) print(f"Model compiled successfully to: {qpc_path}") diff --git a/examples/performance/on_device_sampling.py b/examples/performance/on_device_sampling.py index c34a241c88..4575e4eb33 100644 --- a/examples/performance/on_device_sampling.py +++ b/examples/performance/on_device_sampling.py @@ -56,11 +56,9 @@ def main(args, **kwargs): print("qaic_config:") pprint(qaic_config) - # Load model with On Device Sampler enabled qeff_model = AutoModelForCausalLM.from_pretrained( pretrained_model_name_or_path=args.model_name, continuous_batching=args.full_batch_size is not None, - qaic_config=qaic_config, ) print(f"{args.model_name} optimized for AI 100 \n", qeff_model) @@ -88,6 +86,7 @@ def main(args, **kwargs): mxfp6_matmul=args.mxfp6, mxint8_kv_cache=args.mxint8, num_speculative_tokens=0, + qaic_config=qaic_config, **kwargs, ) print(f"Generated QPC file path: {generated_qpc_path}") diff --git a/examples/performance/speculative_decoding/draft_based.py b/examples/performance/speculative_decoding/draft_based.py index 9e617663ca..4a29d88f85 100644 --- a/examples/performance/speculative_decoding/draft_based.py +++ b/examples/performance/speculative_decoding/draft_based.py @@ -199,9 +199,8 @@ def draft_spec_decode_inference( # export_and_compile tlm and dlm continuous_batching = full_batch_size is not None if target_model_session is None: - target_model = AutoModelForCausalLM.from_pretrained( - target_model_name, continuous_batching=continuous_batching, qaic_config={"speculative_model_type": "target"} - ) + target_qaic_config = {"speculative_model_type": "target"} + target_model = AutoModelForCausalLM.from_pretrained(target_model_name, continuous_batching=continuous_batching) target_num_devices = len(target_device_group) target_model_qpc_path: str = target_model.compile( num_cores=11, @@ -211,6 +210,7 @@ def draft_spec_decode_inference( aic_enable_depth_first=True, full_batch_size=full_batch_size, num_speculative_tokens=num_speculative_tokens, + qaic_config=target_qaic_config, ) target_model_session = QAICInferenceSession(target_model_qpc_path, device_ids=target_device_group) if draft_model_session is None: diff --git a/examples/performance/speculative_decoding/multi_projection.py b/examples/performance/speculative_decoding/multi_projection.py index bbc7456855..1e0bc0ea28 100644 --- a/examples/performance/speculative_decoding/multi_projection.py +++ b/examples/performance/speculative_decoding/multi_projection.py @@ -370,7 +370,6 @@ def get_session( qeff_model = AutoModelForCausalLM.from_pretrained( pretrained_model_name_or_path, continuous_batching=is_cb, - qaic_config=qaic_config, ) num_devices = len(device_group) model_qpc_path: str = qeff_model.compile( @@ -380,6 +379,7 @@ def get_session( ctx_len=ctx_len, aic_enable_depth_first=True, full_batch_size=full_batch_size, + qaic_config=qaic_config, ) print(f"{model_qpc_path=}") # init qaic session diff --git a/examples/performance/speculative_decoding/prompt_lookup.py b/examples/performance/speculative_decoding/prompt_lookup.py index 4a0b0f8141..feec9190f6 100644 --- a/examples/performance/speculative_decoding/prompt_lookup.py +++ b/examples/performance/speculative_decoding/prompt_lookup.py @@ -263,9 +263,8 @@ def pld_spec_decode_inference( # export_and_compile tlm and dlm continuous_batching = full_batch_size is not None - target_model = AutoModelForCausalLM.from_pretrained( - target_model_name, continuous_batching=continuous_batching, qaic_config={"speculative_model_type": "target"} - ) + target_qaic_config = {"speculative_model_type": "target"} + target_model = AutoModelForCausalLM.from_pretrained(target_model_name, continuous_batching=continuous_batching) num_devices = len(device_group) target_model_qpc_path: str = target_model.compile( @@ -276,6 +275,7 @@ def pld_spec_decode_inference( aic_enable_depth_first=True, full_batch_size=full_batch_size, num_speculative_tokens=decode_ks, + qaic_config=target_qaic_config, ) # init qaic session target_model_session = QAICInferenceSession(target_model_qpc_path, device_ids=device_group) diff --git a/examples/text_generation/run_kimik2.py b/examples/text_generation/run_kimik2.py index 0c0acf3878..923f5318b5 100644 --- a/examples/text_generation/run_kimik2.py +++ b/examples/text_generation/run_kimik2.py @@ -49,7 +49,7 @@ # out = model(**inputs) # predictions = torch.argmax(out.logits, dim=-1) -qeff_model = QEFFAutoModelForCausalLM(model, qaic_config=qaic_config) +qeff_model = QEFFAutoModelForCausalLM(model) qeff_model.transform(ctx_len=CTX_LEN, seq_len=PREFILL_SEQ_LEN, bs=1, num_devices=TS, qaic_config=qaic_config) inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) diff --git a/tests/transformers/models/causal_lm_models/check_causal_models.py b/tests/transformers/models/causal_lm_models/check_causal_models.py index 94a581b4d4..fff527b454 100644 --- a/tests/transformers/models/causal_lm_models/check_causal_models.py +++ b/tests/transformers/models/causal_lm_models/check_causal_models.py @@ -127,7 +127,6 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( is_tlm=is_tlm, pretrained_model_name_or_path=model_name, continuous_batching=continuous_batching, - qaic_config=qaic_config, ) qeff_model.transform( ctx_len=ctx_len, @@ -155,7 +154,7 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( else: pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) - onnx_model_path = qeff_model.export(use_onnx_subfunctions=use_onnx_subfunctions) + onnx_model_path = qeff_model.export(use_onnx_subfunctions=use_onnx_subfunctions, qaic_config=qaic_config) if continuous_batching is False: ort_tokens = api_runner.run_kv_model_on_ort(onnx_model_path, is_tlm=is_tlm) gen_len = ort_tokens.shape[-1] @@ -207,6 +206,7 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( full_batch_size=full_batch_size if continuous_batching else None, use_onnx_subfunctions=use_onnx_subfunctions, onnx_path=onnx_model_path, + qaic_config=qaic_config, **compiler_options, **mdp_compile_kwargs, ) diff --git a/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py b/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py index af8c3b70f0..c6f9b4a689 100644 --- a/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py +++ b/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py @@ -79,7 +79,7 @@ def check_causal_lm_pytorch_vs_kv_vs_ai100( is_tlm = False if num_speculative_tokens is None else True qeff_model = QEFFAutoModelForCausalLM( - copy.deepcopy(model_hf), is_tlm=is_tlm, pretrained_model_name_or_path=model_name, qaic_config=qaic_config + copy.deepcopy(model_hf), is_tlm=is_tlm, pretrained_model_name_or_path=model_name ) pytorch_kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) @@ -88,7 +88,7 @@ def check_causal_lm_pytorch_vs_kv_vs_ai100( assert (pytorch_hf_tokens == pytorch_kv_tokens).all(), ( "Tokens don't match for HF PyTorch model output and KV PyTorch model output" ) - qeff_model.export() + qeff_model.export(qaic_config=qaic_config) qpc_path = qeff_model.compile( prefill_seq_len=prompt_len, ctx_len=ctx_len, @@ -100,6 +100,7 @@ def check_causal_lm_pytorch_vs_kv_vs_ai100( prefill_only=prefill_only, enable_qnn=enable_qnn, qnn_config=qnn_config, + qaic_config=qaic_config, ) exec_info = qeff_model.generate(tokenizer, prompts=Constants.INPUT_STR) gen_len = pytorch_kv_tokens.shape[-1] diff --git a/tests/transformers/models/embedding_models/test_qwen3vl_embedding_mad.py b/tests/transformers/models/embedding_models/test_qwen3vl_embedding_mad.py index 0e3819844b..37fa9fc6d4 100644 --- a/tests/transformers/models/embedding_models/test_qwen3vl_embedding_mad.py +++ b/tests/transformers/models/embedding_models/test_qwen3vl_embedding_mad.py @@ -108,7 +108,6 @@ def test_qwen3_vl_embedding_cpu_vs_ai100_mad_parity(model_name): kv_offload=True, trust_remote_code=True, config=qeff_config, - qaic_config={"export_embedding": True}, ) embedder = QEffQwen3VLEmbedder( @@ -131,6 +130,7 @@ def test_qwen3_vl_embedding_cpu_vs_ai100_mad_parity(model_name): num_devices=1, num_cores=16, mxfp6_matmul=False, + qaic_config={"export_embedding": True}, ) cpu_embeddings = _compute_cpu_embeddings(model_hf=model_hf, embedder=embedder, model_inputs=model_inputs) diff --git a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py index a919a37ead..1ce4df913c 100644 --- a/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py +++ b/tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py @@ -137,7 +137,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( model_name, kv_offload=kv_offload, config=config, - qaic_config=qaic_config, torch_dtype=torch_dtype, ignore_mismatched_sizes=True, ) @@ -147,7 +146,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( model_name, kv_offload=kv_offload, config=config, - qaic_config=qaic_config, torch_dtype=torch_dtype, ignore_mismatched_sizes=True, ) @@ -160,7 +158,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( copy.deepcopy(model_hf), kv_offload=kv_offload, config=model_hf.config, - qaic_config=qaic_config, torch_dtype=torch_dtype, ignore_mismatched_sizes=True, ) diff --git a/tests/transformers/subfunction/test_causal_lm_blocking_subfunction.py b/tests/transformers/subfunction/test_causal_lm_blocking_subfunction.py index 267f39e400..2586003837 100644 --- a/tests/transformers/subfunction/test_causal_lm_blocking_subfunction.py +++ b/tests/transformers/subfunction/test_causal_lm_blocking_subfunction.py @@ -39,7 +39,7 @@ def check_blockedKV_onnx_function_count_with_subfunction( """ # Export with subfunctions, NO blocking model_no_block = load_hf_causal_lm_model(model_name, num_hidden_layers=n_layer, config=config) - qeff_no_block = QEFFAutoModelForCausalLM(model_no_block, pretrained_model_name_or_path=model_name, qaic_config=None) + qeff_no_block = QEFFAutoModelForCausalLM(model_no_block, pretrained_model_name_or_path=model_name) qeff_no_block.export(use_onnx_subfunctions=True, offload_pt_weights=False) onnx_no_block = onnx.load(qeff_no_block.onnx_path, load_external_data=False) num_functions_no_block = len(onnx_no_block.functions) @@ -49,10 +49,8 @@ def check_blockedKV_onnx_function_count_with_subfunction( qaic_config = dict(blocking_mode="kv", num_kv_blocks=NUM_KV_BLOCKS) model_kv_block = load_hf_causal_lm_model(model_name, num_hidden_layers=n_layer, config=config) - qeff_kv_block = QEFFAutoModelForCausalLM( - model_kv_block, pretrained_model_name_or_path=model_name, qaic_config=qaic_config - ) - qeff_kv_block.export(use_onnx_subfunctions=True, offload_pt_weights=False) + qeff_kv_block = QEFFAutoModelForCausalLM(model_kv_block, pretrained_model_name_or_path=model_name) + qeff_kv_block.export(use_onnx_subfunctions=True, offload_pt_weights=False, qaic_config=qaic_config) onnx_kv_block = onnx.load(qeff_kv_block.onnx_path, load_external_data=False) num_functions_kv_block = len(onnx_kv_block.functions) diff --git a/tests/transformers/test_pytorch_transforms.py b/tests/transformers/test_pytorch_transforms.py index 442609e7ad..738de4d4cc 100644 --- a/tests/transformers/test_pytorch_transforms.py +++ b/tests/transformers/test_pytorch_transforms.py @@ -190,7 +190,7 @@ def run_kv_cache_transform_and_test( qaic_config = None if "num_logits_to_keep" in qaic_model_inputs: qaic_config = dict(speculative_model_type="target") - qeff_model = QEFFAutoModelForCausalLM(hf_model, qaic_config=qaic_config) + qeff_model = QEFFAutoModelForCausalLM(hf_model) ctx_len = qaic_model_inputs["past_key_values"][0][0].shape[2] qeff_model.transform(ctx_len=ctx_len, seq_len=input_len, bs=input_ids.shape[0], qaic_config=qaic_config) hf_model = qeff_model.model diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 83670e0566..c5ad8a778b 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -1349,7 +1349,8 @@ def test_repeat_kv_quickcheck_hf_qeff_ort_parity(tmp_path): ) torch.manual_seed(0) model_hf = AutoModelForCausalLM.from_config(config, **MODEL_KWARGS).eval() - qeff_model = QEFFAutoModelForCausalLM(deepcopy(model_hf), qaic_config={"replicate_kv_heads": True}) + qaic_config = {"replicate_kv_heads": True} + qeff_model = QEFFAutoModelForCausalLM(deepcopy(model_hf)) input_ids = torch.arange(1, 5, dtype=torch.int64).view(1, 4) position_ids = torch.arange(4, dtype=torch.int64).view(1, 4) @@ -1365,7 +1366,7 @@ def test_repeat_kv_quickcheck_hf_qeff_ort_parity(tmp_path): with torch.no_grad(): hf_logits = model_hf(input_ids=input_ids, position_ids=position_ids).logits[:, -1:, :].detach().numpy() - qeff_model.transform(ctx_len=8, seq_len=4, bs=1, num_devices=4, qaic_config=qeff_model.model.qaic_config) + qeff_model.transform(ctx_len=8, seq_len=4, bs=1, num_devices=4, qaic_config=qaic_config) with torch.no_grad(): qeff_logits = qeff_model.model(**inputs).logits.detach().numpy() diff --git a/tests/unit_test/models/test_modeling_auto_cpu.py b/tests/unit_test/models/test_modeling_auto_cpu.py index b2b58efe4e..f4905eea6f 100644 --- a/tests/unit_test/models/test_modeling_auto_cpu.py +++ b/tests/unit_test/models/test_modeling_auto_cpu.py @@ -24,6 +24,7 @@ Run with: pytest tests/unit_test/models/test_modeling_auto_cpu.py -n auto -v """ +import inspect import logging import os from unittest.mock import MagicMock @@ -49,6 +50,7 @@ QEFFAutoModel, QEFFAutoModelForCausalLM, QEFFAutoModelForCTC, + QEFFAutoModelForImageTextToText, QEFFAutoModelForSequenceClassification, QEFFAutoModelForSpeechSeq2Seq, ) @@ -274,6 +276,40 @@ def test_num_layers_set_correctly(self): qeff = QEFFAutoModelForCausalLM(model) assert qeff.num_layers == 2 + def test_from_pretrained_signature_omits_qaic_config(self): + """from_pretrained no longer exposes qaic_config; compile owns it.""" + causal_from_pretrained_signature = inspect.signature(QEFFAutoModelForCausalLM.from_pretrained) + causal_compile_signature = inspect.signature(QEFFAutoModelForCausalLM.compile) + vlm_from_pretrained_signature = inspect.signature(QEFFAutoModelForImageTextToText.from_pretrained) + + assert "qaic_config" not in causal_from_pretrained_signature.parameters + assert "qaic_config" in causal_compile_signature.parameters + assert "qaic_config" not in vlm_from_pretrained_signature.parameters + + def test_from_pretrained_qaic_config_warns_and_is_not_passed_to_hf(self, monkeypatch): + """Legacy from_pretrained(qaic_config=...) warns and keeps HF loading clean.""" + model, _ = make_tiny_llama() + captured_kwargs = {} + + def fake_from_pretrained(*args, **kwargs): + captured_kwargs.update(kwargs) + return model + + monkeypatch.setattr( + QEFFAutoModelForCausalLM._hf_auto_class, + "from_pretrained", + staticmethod(fake_from_pretrained), + ) + + with pytest.warns(DeprecationWarning, match=r"Pass `qaic_config` to `compile\(\)`"): + qeff = QEFFAutoModelForCausalLM.from_pretrained( + "local-model", + qaic_config={"speculative_model_type": "target"}, + ) + + assert "qaic_config" not in captured_kwargs + assert qeff.model.qaic_config["speculative_model_type"] == "target" + def test_init_raises_type_error_for_non_causal_lm(self): """__init__ raises TypeError when model is not a CausalLM or LMHeadModel.""" model, cfg = make_tiny_bert() @@ -1203,7 +1239,7 @@ def test_compile_list_produces_correct_spec_count(self): from unittest.mock import patch model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) captured = {} with patch.object( @@ -1213,19 +1249,27 @@ def test_compile_list_produces_correct_spec_count(self): captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" ), ): - qeff.compile(prefill_seq_len=32, ctx_len=128, num_speculative_tokens=[0, 3]) + qeff.compile( + prefill_seq_len=32, + ctx_len=128, + num_speculative_tokens=[0, 3], + qaic_config={"speculative_model_type": "target"}, + ) assert captured.get("specializations") is not None, "_compile was not reached" + assert qeff.is_tlm is True + assert qeff.model.qaic_config["speculative_model_type"] == "target" specs = captured["specializations"] decode_specs = [s for s in specs if s.get("seq_len", 0) != 32] assert len(decode_specs) == 2, f"Expected 2 decode specs, got {len(decode_specs)}: {specs}" + assert [s["seq_len"] for s in decode_specs] == [1, 4] def test_compile_deduplication(self): """compile(num_speculative_tokens=[3, 3, 3]) → only one decode spec for K=3.""" from unittest.mock import patch model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) captured = {} with patch.object( @@ -1235,7 +1279,12 @@ def test_compile_deduplication(self): captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" ), ): - qeff.compile(prefill_seq_len=32, ctx_len=128, num_speculative_tokens=[3, 3, 3]) + qeff.compile( + prefill_seq_len=32, + ctx_len=128, + num_speculative_tokens=[3, 3, 3], + qaic_config={"speculative_model_type": "target"}, + ) assert captured.get("specializations") is not None, "_compile was not reached" specs = captured["specializations"] @@ -1248,7 +1297,7 @@ def test_compile_sorting(self): from unittest.mock import patch model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) captured = {} with patch.object( @@ -1258,7 +1307,12 @@ def test_compile_sorting(self): captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" ), ): - qeff.compile(prefill_seq_len=32, ctx_len=128, num_speculative_tokens=[3, 1, 2]) + qeff.compile( + prefill_seq_len=32, + ctx_len=128, + num_speculative_tokens=[3, 1, 2], + qaic_config={"speculative_model_type": "target"}, + ) assert captured.get("specializations") is not None, "_compile was not reached" specs = captured["specializations"] @@ -1272,7 +1326,7 @@ def test_compile_int_backward_compat(self): from unittest.mock import patch model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) captured = {} with patch.object( @@ -1282,7 +1336,12 @@ def test_compile_int_backward_compat(self): captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" ), ): - qeff.compile(prefill_seq_len=32, ctx_len=128, num_speculative_tokens=3) + qeff.compile( + prefill_seq_len=32, + ctx_len=128, + num_speculative_tokens=3, + qaic_config={"speculative_model_type": "target"}, + ) assert captured.get("specializations") is not None, "_compile was not reached" specs = captured["specializations"] @@ -1295,7 +1354,7 @@ def test_compile_int_zero_backward_compat(self): from unittest.mock import patch model, _ = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) captured = {} with patch.object( @@ -1305,7 +1364,12 @@ def test_compile_int_zero_backward_compat(self): captured.update({"specializations": kw.get("specializations")}) or "/fake/qpc" ), ): - qeff.compile(prefill_seq_len=32, ctx_len=128, num_speculative_tokens=0) + qeff.compile( + prefill_seq_len=32, + ctx_len=128, + num_speculative_tokens=0, + qaic_config={"speculative_model_type": "target"}, + ) assert captured.get("specializations") is not None, "_compile was not reached" specs = captured["specializations"] diff --git a/tests/unit_test/transforms/test_speculative_decoding.py b/tests/unit_test/transforms/test_speculative_decoding.py index 3c33fbaec9..b0676bdc62 100644 --- a/tests/unit_test/transforms/test_speculative_decoding.py +++ b/tests/unit_test/transforms/test_speculative_decoding.py @@ -618,11 +618,11 @@ def test_tlm_onnx_has_num_logits_to_keep_input(self, tmp_export_dir): from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM model, cfg = make_tiny_llama() - qeff_model = QEFFAutoModelForCausalLM( - model, + qeff_model = QEFFAutoModelForCausalLM(model) + onnx_path = qeff_model.export( + export_dir=str(tmp_export_dir), qaic_config={"speculative_model_type": "target"}, ) - onnx_path = qeff_model.export(export_dir=str(tmp_export_dir)) onnx_model = onnx.load(str(onnx_path)) input_names = [inp.name for inp in onnx_model.graph.input] @@ -639,11 +639,11 @@ def test_tlm_onnx_logits_output_is_present(self, tmp_export_dir): from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM model, cfg = make_tiny_llama() - qeff_model = QEFFAutoModelForCausalLM( - model, + qeff_model = QEFFAutoModelForCausalLM(model) + onnx_path = qeff_model.export( + export_dir=str(tmp_export_dir), qaic_config={"speculative_model_type": "target"}, ) - onnx_path = qeff_model.export(export_dir=str(tmp_export_dir)) onnx_model = onnx.load(str(onnx_path)) output_names = [out.name for out in onnx_model.graph.output] diff --git a/tests/unit_test/transforms/test_transform_accuracy.py b/tests/unit_test/transforms/test_transform_accuracy.py index fe061b42f0..2486e26a6e 100644 --- a/tests/unit_test/transforms/test_transform_accuracy.py +++ b/tests/unit_test/transforms/test_transform_accuracy.py @@ -256,7 +256,7 @@ def _tiny_llama_qeff(num_attention_heads=4, num_key_value_heads=2): num_key_value_heads=num_key_value_heads, max_position_embeddings=64, ) - return QEFFAutoModelForCausalLM(AutoModelForCausalLM.from_config(cfg), qaic_config={}) + return QEFFAutoModelForCausalLM(AutoModelForCausalLM.from_config(cfg)) def test_repeat_kv_dummy_causal_config(self): qeff_model = self._tiny_llama_qeff() @@ -313,7 +313,7 @@ def test_repeat_kv_dummy_vlm_config(self): vision_feature_layer=-1, ) model_hf = AutoModelForImageTextToText.from_config(cfg) - qeff_model = QEFFAutoModelForImageTextToText(copy.deepcopy(model_hf), kv_offload=False, qaic_config={}) + qeff_model = QEFFAutoModelForImageTextToText(copy.deepcopy(model_hf), kv_offload=False) text_model_before = get_text_model(qeff_model.model) attn_before = get_attention_module(text_model_before.layers[0]) @@ -410,7 +410,7 @@ def test_repeat_kv_skips_encoder_wrapper_without_config(self): vision_feature_layer=-1, ) model_hf = AutoModelForImageTextToText.from_config(cfg) - qeff_model = QEFFAutoModelForImageTextToText(copy.deepcopy(model_hf), kv_offload=True, qaic_config={}) + qeff_model = QEFFAutoModelForImageTextToText(copy.deepcopy(model_hf), kv_offload=True) qeff_model.vision_model.transform(ctx_len=64, seq_len=8, bs=1, qaic_config={"replicate_kv_heads": True}) assert not hasattr(qeff_model.vision_model.model, "config") assert qeff_model.vision_model.hash_params["num_replicate_kv_heads"] == 1 diff --git a/tests/unit_test/utils/test_error_handling.py b/tests/unit_test/utils/test_error_handling.py index 7663457e2a..45fb80a8f3 100644 --- a/tests/unit_test/utils/test_error_handling.py +++ b/tests/unit_test/utils/test_error_handling.py @@ -174,9 +174,10 @@ class TestCheckNumSpeculativeTokensErrorPaths: """check_and_get_num_speculative_tokens must raise for invalid TLM configurations.""" def test_tlm_without_num_speculative_tokens_raises_type_error(self): - """TLM model without num_speculative_tokens must raise TypeError.""" + """TLM model without num_speculative_tokens must raise TypeError after compile-time activation.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) + qeff._activate_qaic_config({"speculative_model_type": "target"}) assert qeff.is_tlm is True with pytest.raises(TypeError, match="num_speculative_tokens"): qeff.check_and_get_num_speculative_tokens(num_speculative_tokens=None, prefill_seq_len=32) @@ -184,7 +185,8 @@ def test_tlm_without_num_speculative_tokens_raises_type_error(self): def test_tlm_prefill_seq_len_too_short_raises_value_error(self): """TLM with prefill_seq_len < num_speculative_tokens+1 must raise ValueError.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) + qeff._activate_qaic_config({"speculative_model_type": "target"}) assert qeff.is_tlm is True # num_speculative_tokens=5, so need prefill_seq_len >= 6 with pytest.raises(ValueError, match="sequence length"): @@ -196,7 +198,8 @@ def test_tlm_prefill_seq_len_too_short_raises_value_error(self): def test_tlm_valid_num_speculative_tokens_does_not_raise(self): """TLM with valid num_speculative_tokens must not raise.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) + qeff._activate_qaic_config({"speculative_model_type": "target"}) result = qeff.check_and_get_num_speculative_tokens(num_speculative_tokens=3, prefill_seq_len=32) assert result == 3 @@ -329,32 +332,31 @@ def test_is_tlm_false_without_config(self): qeff = QEFFAutoModelForCausalLM(model) assert qeff.is_tlm is False - def test_is_tlm_false_with_empty_config(self): - """is_tlm must be False when qaic_config has no speculative_model_type.""" + def test_is_tlm_false_with_empty_compile_config(self): + """is_tlm must be False when compile-time qaic_config has no speculative_model_type.""" model = make_tiny_gpt2() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={}) + qeff = QEFFAutoModelForCausalLM(model) + qeff._activate_qaic_config({}) assert qeff.is_tlm is False - def test_is_tlm_true_with_target_type(self): - """is_tlm must be True when speculative_model_type='target'.""" + def test_is_tlm_true_with_compile_target_type(self): + """is_tlm must be True when compile-time speculative_model_type='target'.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "target"}) + qeff = QEFFAutoModelForCausalLM(model) + qeff._activate_qaic_config({"speculative_model_type": "target"}) assert qeff.is_tlm is True def test_turbo_type_requires_pretrained_model_name(self): """speculative_model_type='turbo' without pretrained_model_name_or_path must raise KeyError.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM(model, qaic_config={"speculative_model_type": "turbo"}) + qeff = QEFFAutoModelForCausalLM(model) with pytest.raises(KeyError, match="pretrained_model_name_or_path"): - qeff.transform() + qeff.transform(qaic_config={"speculative_model_type": "turbo"}) def test_cb_and_tlm_together_model_is_tlm(self): """continuous_batching=True with TLM: model must still be recognized as TLM.""" model = make_tiny_llama() - qeff = QEFFAutoModelForCausalLM( - model, - continuous_batching=True, - qaic_config={"speculative_model_type": "target"}, - ) + qeff = QEFFAutoModelForCausalLM(model, continuous_batching=True) + qeff._activate_qaic_config({"speculative_model_type": "target"}) # The model should be recognized as TLM regardless of CB flag assert qeff.is_tlm is True diff --git a/tests/weight_free/test_ccl.py b/tests/weight_free/test_ccl.py index 3fe116acd8..b305c5b207 100644 --- a/tests/weight_free/test_ccl.py +++ b/tests/weight_free/test_ccl.py @@ -48,9 +48,7 @@ def test_weight_free_ccl_compile_and_generate(model_type, model_id, tmp_export_d try: config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) config.num_hidden_layers = 2 - qeff_model = QEFFAutoModelForCausalLM.from_pretrained( - model_id, config=config, weight_free=True, qaic_config={"ccl_enabled": True} - ) + qeff_model = QEFFAutoModelForCausalLM.from_pretrained(model_id, config=config, weight_free=True) tokenizer = load_tokenizer(model_id) except Exception as exc: skip_on_model_fetch_error(exc, model_id) @@ -72,6 +70,7 @@ def test_weight_free_ccl_compile_and_generate(model_type, model_id, tmp_export_d num_cores=16, batch_size=BATCH_SIZE, use_onnx_subfunctions=True, + qaic_config={"ccl_enabled": True}, ) output = qeff_model.generate( tokenizer=tokenizer, @@ -91,6 +90,7 @@ def test_weight_free_ccl_compile_and_generate(model_type, model_id, tmp_export_d num_cores=16, batch_size=BATCH_SIZE, use_onnx_subfunctions=True, + qaic_config={"ccl_enabled": True}, ) output = qeff_model.generate( tokenizer=tokenizer,