diff --git a/QEfficient/__init__.py b/QEfficient/__init__.py index 393ee7d4ad..fb85183144 100755 --- a/QEfficient/__init__.py +++ b/QEfficient/__init__.py @@ -19,6 +19,7 @@ import warnings # noqa: I001 import transformers import transformers.utils as transformers_utils +from transformers.utils import import_utils as hf_import_utils try: from transformers import HybridCache as _TransformersHybridCache # noqa: F401 @@ -117,3 +118,23 @@ def check_qaic_sdk(): if not check_qaic_sdk(): logger.warning("QAIC SDK is not installed, eager mode features won't be available!") + + +def ensure_torch_fx_import_compatibility(): + if hasattr(hf_import_utils, "is_torch_fx_available"): + return + + def _is_torch_fx_available() -> bool: + if not hf_import_utils.is_torch_available(): + return False + try: + import torch.fx # noqa: F401 + + return True + except Exception: + return False + + hf_import_utils.is_torch_fx_available = _is_torch_fx_available + + +ensure_torch_fx_import_compatibility() diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 6be24bee44..3d89596c07 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -14,7 +14,7 @@ import warnings from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Type, Union import onnx import torch @@ -117,6 +117,7 @@ class QEFFBaseModel(ABC): _layerwise_active = False _pytorch_transforms: List[PytorchTransform] _onnx_transforms = [BaseOnnxTransform] + _checkpoint_transforms: List[Type] = [] def _transform_names(self) -> List[str]: return [x.__name__ for x in self._pytorch_transforms + self._onnx_transforms] @@ -160,6 +161,7 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None: self.onnx_path: Optional[str] = None self.qpc_path: Optional[str] = None self.qpc_session: Optional[QAICInferenceSession] = None + self.weight_spec_path: Optional[str] = None self.model_architecture = ( (arch := getattr(self.model.config, "architectures", None)) and len(arch) > 0 and arch[0] ) or None @@ -432,6 +434,78 @@ def _export_via_dynamo( else: os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = prev_invoke_fallback + def _export_via_weightfree( + self, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_axes: Dict, + export_kwargs: Dict, + onnx_transform_kwargs: Optional[Dict] = None, + ): + """Export through the weight-free dynamo path with checkpoint-backed weights.""" + from QEfficient.customop.dynamo_ops import DYNAMO_CUSTOM_OP_TABLE + from QEfficient.exporter.weight_free.core import export_weight_free_onnx + from QEfficient.utils.export_utils import convert_dynamic_axes_to_dynamic_shapes + + model_config = getattr(self.model, "config", None) + dynamic_shapes = convert_dynamic_axes_to_dynamic_shapes(dynamic_axes, model_config) + + sig_keys = list(inspect.signature(self.model.forward).parameters.keys()) + sig_key_set = set(sig_keys) + ordered_inputs, ordered_shapes = {}, {} + for key in sig_keys: + if key in example_inputs: + ordered_inputs[key] = example_inputs[key] + if key in dynamic_shapes: + ordered_shapes[key] = dynamic_shapes[key] + example_inputs = { + **ordered_inputs, + **{key: value for key, value in example_inputs.items() if key not in sig_key_set}, + } + dynamic_shapes = { + **ordered_shapes, + **{key: value for key, value in dynamic_shapes.items() if key not in sig_key_set}, + } + + wf_export_kwargs = dict(export_kwargs) + wf_export_kwargs.setdefault("report", False) + wf_export_kwargs.setdefault("optimize", False) + wf_export_kwargs["dynamo"] = True + wf_export_kwargs["opset_version"] = constants.ONNX_DYNAMO_EXPORT_OPSET + wf_export_kwargs["custom_translation_table"] = { + **(wf_export_kwargs.pop("custom_translation_table", None) or {}), + **DYNAMO_CUSTOM_OP_TABLE, + } + + export_func = export_weight_free_onnx + if self.model.__class__.__name__ in {"QEffKimiK25DecoderWrapper", "QEffKimiK25EncoderWrapper"}: + from QEfficient.exporter.weight_free.core import export_loaded_model_weight_free_onnx + + export_func = export_loaded_model_weight_free_onnx + + prev_invoke_fallback = os.environ.get("TORCH_INVOKE_ALLOW_CREATE_FALLBACK") + os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = "1" + try: + _, updated_onnx_transform_kwargs, cleanup = export_func( + qeff_model=self, + tmp_onnx_path=tmp_onnx_path, + example_inputs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_shapes=dynamic_shapes, + export_kwargs=wf_export_kwargs, + onnx_transform_kwargs=onnx_transform_kwargs or {}, + ) + finally: + if prev_invoke_fallback is None: + os.environ.pop("TORCH_INVOKE_ALLOW_CREATE_FALLBACK", None) + else: + os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = prev_invoke_fallback + + return updated_onnx_transform_kwargs, cleanup + @export_wrapper def _export( self, @@ -444,6 +518,7 @@ def _export( prefill_only: Optional[bool] = False, dynamo: bool = False, dynamic_shapes: Optional[Dict[str, Dict[int, Any]]] = None, + use_weight_free_export: bool = False, **export_kwargs, ) -> str: """ @@ -476,14 +551,21 @@ def _export( # TODO: Hack for retain_full_kv, handle this outside export_kwargs.pop("retain_full_kv", None) onnx_path = export_dir / f"{self.model_name}.onnx" + weight_spec_path = onnx_path.with_name("weight_spec.json") # Return early if ONNX already exists if onnx_path.is_file(): self.onnx_path = onnx_path + if weight_spec_path.is_file(): + self.weight_spec_path = str(weight_spec_path) return onnx_path # check if the model is in meta state or weights are offloaded - self._model_offloaded_check() + if not use_weight_free_export: + self._model_offloaded_check() + + if use_weight_free_export and not dynamo: + raise NotImplementedError("Weight-free export requires dynamo=True.") export_dir.mkdir(parents=True, exist_ok=True) @@ -556,8 +638,28 @@ def _resolve_pkv_names(layer_idx, layer_state): dynamic_axes = {rename_map.get(k, k): v for k, v in dynamic_axes.items()} input_names = aligned_input_names + cleanup_fn = None try: - if dynamo: + if use_weight_free_export: + tmp_onnx_dir = export_dir / "onnx_weightfree_tmp" + tmp_onnx_dir.mkdir(parents=True, exist_ok=True) + tmp_onnx_path = tmp_onnx_dir / onnx_path.name + onnx_transform_kwargs, cleanup_fn = self._export_via_weightfree( + tmp_onnx_path=tmp_onnx_path, + example_inputs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + export_kwargs=export_kwargs, + onnx_transform_kwargs=onnx_transform_kwargs, + ) + shutil.move(str(tmp_onnx_path), str(onnx_path)) + tmp_weight_spec = tmp_onnx_dir / "weight_spec.json" + if tmp_weight_spec.exists(): + shutil.move(str(tmp_weight_spec), str(weight_spec_path)) + self.weight_spec_path = str(weight_spec_path) + shutil.rmtree(tmp_onnx_dir, ignore_errors=True) + elif dynamo: self._export_via_dynamo( onnx_path, example_inputs, @@ -576,7 +678,8 @@ def _resolve_pkv_names(layer_idx, layer_state): export_kwargs, ) logger.info("PyTorch export successful") - self._offload_model_weights(offload_pt_weights) + if not use_weight_free_export: + self._offload_model_weights(offload_pt_weights) model = onnx.load(onnx_path, load_external_data=False) needs_external_tensor_data = any( @@ -585,13 +688,18 @@ def _resolve_pkv_names(layer_idx, layer_state): transform_kwargs = { "onnx_base_dir": str(export_dir) if needs_external_tensor_data else None, "model_name": self.model_name, - "dynamic_axes": None if dynamo else dynamic_axes, # dynamo uses dynamic_shapes, not axes - "onnx_export_opset": constants.get_onnx_export_opset(dynamo), + "dynamic_axes": None if (dynamo or use_weight_free_export) else dynamic_axes, + "onnx_export_opset": constants.get_onnx_export_opset(dynamo or use_weight_free_export), } if onnx_transform_kwargs is not None: transform_kwargs.update(onnx_transform_kwargs) - onnx_transforms = OnnxTransformPipeline(transforms=self._onnx_transforms) + active_transforms = [ + transform + for transform in self._onnx_transforms + if not (use_weight_free_export and transform is SplitTensorsTransform) + ] + onnx_transforms = OnnxTransformPipeline(transforms=active_transforms) model, transformed = onnx_transforms.apply(model, **transform_kwargs) # Keep this strictly layerwise-scoped so regular non-layerwise export @@ -603,6 +711,15 @@ def _resolve_pkv_names(layer_idx, layer_state): model.metadata_props.append( onnx.StringStringEntryProto(key="qeff_transforms", value=",".join(self._transform_names())) ) + if use_weight_free_export and self.weight_spec_path is not None: + import json as _json + + from QEfficient.exporter.weight_free.core import _upsert_metadata_prop + + weight_spec_json = _json.dumps( + load_json(Path(self.weight_spec_path)), separators=(",", ":"), sort_keys=True + ) + _upsert_metadata_prop(model, "com.qti.aisw.extdata", weight_spec_json) logger.info("ONNX transforms applied") onnx_path_tmp = onnx_path.with_suffix(onnx_path.suffix + ".tmp") @@ -615,8 +732,19 @@ def _resolve_pkv_names(layer_idx, layer_state): except Exception as e: logger.error(f"ONNX export or transforms failed: {e}") raise e + finally: + if cleanup_fn is not None: + cleanup_fn() self.onnx_path = onnx_path + if use_weight_free_export and self.weight_spec_path is not None: + from QEfficient.exporter.weight_free.spec import load_weight_spec + + spec = load_weight_spec(Path(self.weight_spec_path)) + prepared_out = Path(spec.model_id) + symlink = onnx_path.parent / prepared_out.name + if prepared_out.exists() and not symlink.exists(): + symlink.symlink_to(prepared_out) return onnx_path def get_onnx_path( @@ -631,6 +759,7 @@ def get_onnx_path( qaic_config: Optional[dict] = None, moe_prefill_packed_chunk_size: Optional[int] = None, kv_cache_prefix: Optional[str] = None, + use_weight_free_export: bool = False, **compiler_options, ): kwargs = { @@ -638,6 +767,7 @@ def get_onnx_path( "use_onnx_subfunctions": use_onnx_subfunctions, "dynamo": dynamo, "retain_full_kv": retain_full_kv, + "use_weight_free_export": use_weight_free_export, } layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) if layerwise_cache_probe: @@ -1009,6 +1139,7 @@ def _compile( layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) moe_prefill_packed_chunk_size = compiler_options.pop("moe_prefill_packed_chunk_size", None) + use_weight_free_export = compiler_options.pop("use_weight_free_export", False) for removed_option in ("compile_only", "compile-only"): if removed_option in compiler_options: @@ -1040,6 +1171,7 @@ def _compile( moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) if QEFFBaseModel._layerwise_active: diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index eb27051b9d..c35de5daf4 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -44,9 +44,8 @@ CtxScatterFuncCB, CtxScatterFuncCB3D, ) - -# from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func from QEfficient.customop.onnxscript_utils import get_onnxscript_func +from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func, update_cast_to_uint4_output_types from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc from QEfficient.utils import constants from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT @@ -112,7 +111,7 @@ class CustomOpTransform(BaseOnnxTransform): "CtxGatherFuncBlockedKVCB": (CtxGatherFuncBlockedKVCB, CtxGatherBlockedKVCB), "CtxScatterFuncCB": (CtxScatterFuncCB, CtxScatterCB), "CtxGatherFuncCB": (CtxGatherFuncCB, CtxGatherCB), - # "CastToUInt4": (CastToUInt4Func, CastToUInt4), + "CastToUInt4": (CastToUInt4Func, CastToUInt4), } @classmethod @@ -598,7 +597,7 @@ def apply( **kwargs, ) -> Tuple[ModelProto, bool]: if not self.transforms: - return model, False + return model, update_cast_to_uint4_output_types(model) # Same logic as before, but replace `transforms` with `self.transforms` mapping: Dict[str, Tuple[TensorProto, str]] = {} @@ -649,6 +648,8 @@ def _set_external_data(tensor, file_name): model, onnx_export_opset=kwargs.get("onnx_export_opset", constants.ONNX_LEGACY_EXPORT_OPSET) ) + cast_to_uint4_types_updated = update_cast_to_uint4_output_types(model) + if RenameFunctionOutputsTransform in requested: applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply( model, layer_idx=kwargs.get("layer_idx", 0) @@ -666,4 +667,4 @@ def _set_external_data(tensor, file_name): for t, done in applied.items(): logger.info(f"Transform '{t.__name__}' applied={done}") - return model, any(applied.values()) + return model, any(applied.values()) or cast_to_uint4_types_updated diff --git a/QEfficient/base/pytorch_transforms.py b/QEfficient/base/pytorch_transforms.py index 812177eac2..a716805d36 100644 --- a/QEfficient/base/pytorch_transforms.py +++ b/QEfficient/base/pytorch_transforms.py @@ -5,7 +5,7 @@ # # ---------------------------------------------------------------------------- from types import MethodType -from typing import Callable, Dict, Tuple, Type +from typing import Callable, Dict, Optional, Tuple, Type from torch import nn @@ -97,6 +97,7 @@ class ModuleMutatorTransform(PytorchTransform): """ _match_class: nn.Module + _match_string: Optional[str] = None @classmethod def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]: @@ -135,7 +136,14 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]: repl_method_map := cls._match_string_replace_method.get(module.__class__.__name__) ): for orig_method_name, mapped_method in repl_method_map.items(): - setattr(module, orig_method_name, MethodType(mapped_method, module)) + parts = orig_method_name.split(".") + if len(parts) > 1: + target = module + for part in parts[:-1]: + target = getattr(target, part) + setattr(target, parts[-1], MethodType(mapped_method, target)) + else: + setattr(module, orig_method_name, MethodType(mapped_method, module)) if hasattr(module, "__qeff_init__"): module.__qeff_init__() diff --git a/QEfficient/customop/dynamo_ops.py b/QEfficient/customop/dynamo_ops.py index 81318efbbf..3b9a6152aa 100644 --- a/QEfficient/customop/dynamo_ops.py +++ b/QEfficient/customop/dynamo_ops.py @@ -23,6 +23,7 @@ CtxScatterCB3D, ) from QEfficient.customop.onnxscript_utils import get_dynamo_onnxscript_func +from QEfficient.customop.quantization_ops import CastToUInt4 from QEfficient.customop.rms_norm import CustomRMSNorm # noqa: E402 @@ -40,6 +41,21 @@ def _(hidden_states: torch.Tensor, weight: torch.Tensor, epsilon: float) -> torc return torch.empty_like(hidden_states) +@torch.library.custom_op("qefficient::cast_to_uint4", mutates_args=()) +def cast_to_uint4_op(weight_packed: torch.Tensor) -> torch.Tensor: + """Unpack packed uint8 nibbles into uint8 values for eager/fake execution.""" + lower = weight_packed & 0x0F + upper = (weight_packed >> 4) & 0x0F + output_shape = (*weight_packed.shape[:-1], weight_packed.shape[-1] * 2) + return torch.stack([lower, upper], dim=-1).reshape(output_shape) + + +@cast_to_uint4_op.register_fake +def _(weight_packed: torch.Tensor) -> torch.Tensor: + output_shape = (*weight_packed.shape[:-1], weight_packed.shape[-1] * 2) + return torch.empty(output_shape, dtype=weight_packed.dtype, device=weight_packed.device) + + @torch.library.custom_op("qefficient::ctx_scatter", mutates_args=()) def ctx_scatter_op(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: """Custom context scatter operation""" @@ -374,6 +390,7 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> DYNAMO_CUSTOM_OP_TABLE = { torch.ops.qefficient.rms_norm.default: get_dynamo_onnxscript_func(CustomRMSNorm), + torch.ops.qefficient.cast_to_uint4.default: get_dynamo_onnxscript_func(CastToUInt4), torch.ops.qefficient.ctx_scatter.default: get_dynamo_onnxscript_func(CtxScatter), torch.ops.qefficient.ctx_scatter_3d.default: get_dynamo_onnxscript_func(CtxScatter3D), torch.ops.qefficient.ctx_scatter_cb.default: get_dynamo_onnxscript_func(CtxScatterCB), diff --git a/QEfficient/customop/matmulnbits.py b/QEfficient/customop/matmulnbits.py index d8cc0e8f1b..72cf898ad5 100644 --- a/QEfficient/customop/matmulnbits.py +++ b/QEfficient/customop/matmulnbits.py @@ -14,7 +14,10 @@ class QuantLinearTorchFunction(torch.autograd.Function): @staticmethod def symbolic(g, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group_size, in_features, out_features): - input_tuple = (x, qself_qweight, qself_scales, qself_qzeros) + if qself_qzeros is None: + input_tuple = (x, qself_qweight, qself_scales) + else: + input_tuple = (x, qself_qweight, qself_scales, qself_qzeros) input_tuple += (g_idx,) if g_idx is not None else () return g.op( "com.microsoft::MatMulNBits", @@ -28,7 +31,10 @@ def symbolic(g, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group @staticmethod def forward(ctx, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, group_size, in_features, out_features): + if qself_qzeros is None: + qself_qzeros = 2 ^ (bits - 1) if torch.onnx.is_in_onnx_export(): + # For faster export return torch.zeros(x.shape[:-1] + (out_features,), dtype=x.dtype).float() fp_weight = dequantize_blockwise_bits( qself_qweight, qself_scales, qself_qzeros, bits, group_size, g_idx, in_features, out_features @@ -40,8 +46,7 @@ def forward(ctx, x, qself_qweight, qself_scales, qself_qzeros, g_idx, bits, grou def dequantize_blockwise_bits(quant_values, scale, zero_point, bits, group_size, g_idx, rows, cols): if bits != 4: raise ValueError("Only bits=4 is supported for executing quantized model") - if group_size != 128: - raise ValueError("Only group_size=128 is supported for executing quantized model") + expand_quant_value = (quant_values.unsqueeze(-1) >> torch.tensor([[[[0, 4]]]], dtype=torch.int32)) & 0x0F expand_quant_value = expand_quant_value.reshape(*quant_values.shape[:-1], -1) aligned_scale = scale.reshape(*quant_values.shape[:-1], 1) @@ -88,20 +93,20 @@ def __init__(self, bits, group_size, in_features, out_features, bias): q_rows = in_features // self.group_size self.register_buffer( "qweight", - torch.zeros((out_features, q_rows, self.group_size // (8 // bits)), dtype=torch.uint8), + torch.empty((out_features, q_rows, self.group_size // (8 // bits)), dtype=torch.uint8), ) self.register_buffer( "qzeros", - torch.zeros((q_rows + (q_rows & 1)) * (out_features // 8 * self.bits), dtype=torch.uint8), + torch.empty((q_rows + (q_rows & 1)) * (out_features // 8 * self.bits), dtype=torch.uint8), ) self.register_buffer( - "scales", torch.zeros((math.ceil(in_features / self.group_size) * out_features), dtype=torch.float16) + "scales", torch.empty((math.ceil(in_features / self.group_size) * out_features), dtype=torch.float16) ) self.register_buffer( "g_idx", torch.tensor([i // self.group_size for i in range(in_features)], dtype=torch.int32) ) if bias: - self.register_buffer("bias", torch.zeros((out_features), dtype=torch.float16)) + self.register_buffer("bias", torch.empty((out_features), dtype=torch.float16)) else: self.bias = None @@ -183,3 +188,72 @@ def forward(self, inputs): ) out = out + self.bias if self.bias is not None else out return out + + +class QMOE(torch.autograd.Function): + @staticmethod + def symbolic( + g, + x, + router_weights, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + router_probs, + activation_type, + block_size, + expert_weight_bits, + k, + ): + qmoe_out = g.op( + "com.microsoft::QMoE", + x, + router_weights, + router_probs, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + outputs=1, + activation_type_s=activation_type, # <-- _s suffix for string + block_size_i=block_size, + expert_weight_bits_i=expert_weight_bits, + k_i=k, + ) + + # # Create axes=-1 as an explicit int64 constant tensor + # axes = g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + + # # Compute mean of router_probs along the last axis, keepdims for broadcasting + # router_probs_mean = g.op("ReduceMean", router_probs, axes, keepdims_i=1) + + # Multiply qmoe_out with the averaged router_probs + return qmoe_out + # return g.op("Mul", qmoe_out, router_probs_mean)) + + @staticmethod + def forward( + ctx, + x, + router_weights, + fc1_experts_weights, + fc1_scales, + fc2_experts_weights, + fc2_scales, + fc3_experts_weights, + fc3_scales, + router_probs, + activation_type, + block_size, + expert_weight_bits, + k, + ): + # Dummy forward: simulate qmoe_out as zeros_like(x), then apply ReduceMean * Mul + qmoe_out = torch.zeros_like(x) + router_probs_mean = router_probs.mean(dim=-1, keepdim=True) + return qmoe_out * router_probs_mean diff --git a/QEfficient/customop/quantization_ops.py b/QEfficient/customop/quantization_ops.py new file mode 100644 index 0000000000..0c205abfbe --- /dev/null +++ b/QEfficient/customop/quantization_ops.py @@ -0,0 +1,200 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import onnxscript +import torch +from onnx import ModelProto, TensorProto +from onnxscript.onnx_types import UINT4 + +from QEfficient.customop.onnxscript_utils import qeff_custom_op +from QEfficient.customop.utils import select_interface +from QEfficient.utils import constants + +ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET)) + + +@qeff_custom_op("com.qti.aisw.onnx", 1) +def CastToUInt4(weight_packed: onnxscript.UINT8) -> UINT4: + """ + Unpack packed uint8 weights into uint4 values and cast output to UINT4. + Supports N-D input: all leading dimensions are preserved; only the last + dimension (in_features // 2) is doubled to (in_features). + + Input: (..., in_features // 2) UINT8 + Each byte holds two nibbles: byte = (w_y << 4) | (w_x & 0x0F) + Output: (..., in_features) UINT4, values in [0, 15] + + Operations: + w_x = weight_packed % 16 (lower nibble) + w_y = (weight_packed >> 4) % 16 (upper nibble) + stacked = concat([w_x, w_y], axis=-1) after unsqueeze + → (..., in//2, 2) + leading_dims = shape[:-1] + new_shape = [...leading_dims, last_dim * 2] + reshaped = reshape(stacked, new_shape) + output = Cast(reshaped, to=UINT4) + """ + sixteen = ops.CastLike(ops.Constant(value_ints=[16]), weight_packed) + + # Lower nibble: weight_packed & 0x0F = weight_packed % 16 + w_x = ops.Mod(weight_packed, sixteen) + + # Upper nibble: (weight_packed >> 4) & 0x0F + shift = ops.CastLike(ops.Constant(value_ints=[4]), weight_packed) + w_shifted = ops.BitShift(weight_packed, shift, direction="RIGHT") + w_y = ops.Mod(w_shifted, sixteen) + + # Stack along a new last dim → (..., in_features//2, 2) + w_x_unsq = ops.Unsqueeze(w_x, [-1]) + w_y_unsq = ops.Unsqueeze(w_y, [-1]) + stacked = ops.Concat(w_x_unsq, w_y_unsq, axis=-1) + + # N-D aware reshape: preserve all leading dims, double the last dim. + # packed_shape = [d0, d1, ..., last_dim] + packed_shape = ops.Shape(weight_packed) + # All dims except the last: [d0, d1, ...] + leading_dims = ops.Slice(packed_shape, starts=[0], ends=[-1], axes=[0]) + # Last dim only: [last_dim] + last_dim = ops.Slice(packed_shape, starts=[-1], ends=[2147483647], axes=[0]) + # Double the last dim: [last_dim * 2] + last_dim_doubled = ops.Mul(last_dim, ops.Constant(value_ints=[2])) + # New shape: [d0, d1, ..., last_dim * 2] + new_shape = ops.Concat(leading_dims, last_dim_doubled, axis=0) + reshaped = ops.Reshape(stacked, new_shape) + + # Cast to UINT4 — data_type value is version-dependent (21 in ONNX 1.18, 23 in newer) + return ops.Cast(reshaped, to=int(TensorProto.UINT4)) + + +class CastToUInt4Func(torch.autograd.Function): + """ + Custom op: unpacks packed uint8 → uint8 (values 0-15) in PyTorch. + In ONNX the custom op subgraph includes a Cast → UINT4 as its last step. + Supports N-D input: all leading dimensions are preserved. + + PyTorch forward : packed uint8 (..., in//2) → uint8 (..., in), values [0, 15] + ONNX symbolic : emits CastToUInt4 node (com.qti.aisw.onnx) + The subgraph ends with Cast → UINT4. + """ + + @staticmethod + def forward(weight_packed: torch.Tensor) -> torch.Tensor: + w_x = weight_packed & 0x0F # lower nibble, (..., in//2), range [0, 15] + w_y = (weight_packed >> 4) & 0x0F # upper nibble, (..., in//2), range [0, 15] + # New shape: all leading dims unchanged, last dim doubled + new_shape = list(weight_packed.shape[:-1]) + [weight_packed.shape[-1] * 2] + return torch.stack( + [w_x, w_y], dim=-1 + ).reshape( + new_shape + ) # Can't add a cast operation to uint4 here, as its not supported in pytorch; The ONNX export will handle the cast to IINT4 in the symbolic method. + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: + output = g.onnxscript_op(CastToUInt4, weight_packed) + return output + + +def cast_to_uint4(weight_packed: torch.Tensor) -> torch.Tensor: + return select_interface(CastToUInt4Func.apply, torch.ops.qefficient.cast_to_uint4)(weight_packed) + + +def update_cast_to_uint4_output_types(model: ModelProto) -> bool: + """Correct exported CastToUInt4 value-info to logical UINT4. + + PyTorch fake tensors cannot carry UINT4 dtype, so Dynamo annotates the + custom-op node outputs as UINT8 even though the ONNXScript function returns + UINT4. QAIC consumes the graph-level value-info for quantized + dequantization, so keep the node metadata consistent with the custom-op + body. + """ + cast_outputs = { + output_name + for node in model.graph.node + if node.domain == "com.qti.aisw.onnx" and node.op_type == "CastToUInt4" + for output_name in node.output + } + if not cast_outputs: + return False + + transformed = False + for value in list(model.graph.value_info) + list(model.graph.output) + list(model.graph.input): + if value.name not in cast_outputs: + continue + tensor_type = value.type.tensor_type + if tensor_type.elem_type != TensorProto.UINT4: + tensor_type.elem_type = TensorProto.UINT4 + transformed = True + return transformed + + +class DequantizeLinearFunc(torch.autograd.Function): + """ + Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom). + + Symmetric blockwise quantization — no zero_point: + output = x * scale (per block along the last axis) + + Supports N-D input: + weight_unpacked : (..., in_features) — quantized values + scale : (..., num_blocks) — per-block scales + block_size : int — elements per block + + PyTorch forward : expand blockwise scale along last dim, multiply + ONNX symbolic : DequantizeLinear(weight_unpacked, scale, + axis=2, block_size=block_size) + axis=2 for 3D input (2, out_features, in_features). + No zero_point input (symmetric). + """ + + @staticmethod + def forward( + weight_unpacked: torch.Tensor, scale: torch.Tensor, zeros: torch.Tensor, block_size: int + ) -> torch.Tensor: + # Expand per-block scale → per-element scale along last dim + scale_expanded = scale.repeat_interleave(block_size, dim=-1) + zeros_expanded = zeros.repeat_interleave(block_size, dim=-1) + return (weight_unpacked.to(torch.int8) - zeros_expanded.to(torch.int8)) * scale_expanded + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def symbolic( + g: torch.Graph, weight_unpacked: torch.Value, scale: torch.Value, zeros: torch.Value, block_size: int + ) -> torch.Value: + # Standard DequantizeLinear: symmetric (no zero_point), blockwise. + # Input is 3D: (2, out_features, in_features) → axis=2 (last dim). + # DequantizeLinear natively supports batch dimensions. + return g.op( + "DequantizeLinear", + weight_unpacked, + scale, + zeros, + axis_i=2, + block_size_i=block_size, + ) + + +def dequantize_linear( + weight_unpacked: torch.Tensor, scale: torch.Tensor, zeros: torch.Tensor, block_size: int +) -> torch.Tensor: + if torch._dynamo.is_compiling(): + return torch.onnx.ops.symbolic( + "::DequantizeLinear", + (weight_unpacked, scale, zeros), + {"axis": 2, "block_size": block_size}, + dtype=scale.dtype, + shape=weight_unpacked.shape, + version=18, + ) + return DequantizeLinearFunc.apply(weight_unpacked, scale, zeros, block_size) diff --git a/QEfficient/exporter/weight_free/__init__.py b/QEfficient/exporter/weight_free/__init__.py new file mode 100644 index 0000000000..6397204542 --- /dev/null +++ b/QEfficient/exporter/weight_free/__init__.py @@ -0,0 +1,43 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +from QEfficient.exporter.weight_free.core import ( + export_weight_free_onnx as export_weight_free_onnx, +) +from QEfficient.exporter.weight_free.core import ( + load_weight_free_ort_inputs as load_weight_free_ort_inputs, +) +from QEfficient.exporter.weight_free.core import ( + log_weight_free_export as log_weight_free_export, +) +from QEfficient.exporter.weight_free.spec import ( + ExternalDataFile as ExternalDataFile, +) +from QEfficient.exporter.weight_free.spec import ( + load_weight_spec as load_weight_spec, +) +from QEfficient.exporter.weight_free.spec import ( + resolve_weight_spec_path as resolve_weight_spec_path, +) +from QEfficient.exporter.weight_free.spec import ( + save_weight_spec as save_weight_spec, +) +from QEfficient.exporter.weight_free.transforms import ( + BaseCheckpointTransform as BaseCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + CheckpointTransformPipeline as CheckpointTransformPipeline, +) +from QEfficient.exporter.weight_free.transforms import ( + DtypeConversionCheckpointTransform as DtypeConversionCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + MoEExpertStackingCheckpointTransform as MoEExpertStackingCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + MoEFusedExpertSplitCheckpointTransform as MoEFusedExpertSplitCheckpointTransform, +) diff --git a/QEfficient/exporter/weight_free/core.py b/QEfficient/exporter/weight_free/core.py new file mode 100644 index 0000000000..4a8dd71a1d --- /dev/null +++ b/QEfficient/exporter/weight_free/core.py @@ -0,0 +1,720 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import os +import sys +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +import numpy as np +import onnx_ir as ir +import torch +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file +from torch import nn + +from QEfficient.transformers.embeddings.embedding_utils import PooledModel +from QEfficient.transformers.models.pytorch_transforms import PoolingTransform +from QEfficient.utils.export_utils import ( + _cleanup_onnx_subfunctions, + _setup_onnx_subfunctions, + get_decoder_layer_classes_for_export, +) +from QEfficient.utils.logging_utils import logger +from QEfficient.utils.torch_patches import ( + temporarily_disable_nested_compile_regions, + temporarily_enable_nested_compile_regions, +) + +from .spec import ( + ExternalDataFile, + TiedWeightAlias, + WeightSpec, + WeightSpecInput, + WeightSpecLocation, + load_weight_spec, + resolve_weight_spec_path, + save_weight_spec, +) +from .transforms import CheckpointTransformPipeline + +# Memory profiler from scripts/memory_profiling — optional, degrades gracefully +# if the scripts directory is not on the path or matplotlib is missing. +sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "scripts")) +try: + from memory_profiling import QEffMemoryProfiler as _QEffMemoryProfiler + + _HAS_PROFILER = True +except ImportError: + _HAS_PROFILER = False + +# Last checkpoint-prep profiler stats — written after every export call so +# _runner.py can read them without text-parsing stdout. +_last_prep_peak_rss_mb: float = 0.0 +_last_prep_duration_seconds: float = 0.0 +_checkpoint_prep_ran: bool = False + + +def _to_meta(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return torch.empty_like(value, device="meta") + if isinstance(value, tuple): + return tuple(_to_meta(item) for item in value) + if isinstance(value, list): + return [_to_meta(item) for item in value] + if isinstance(value, dict): + return {key: _to_meta(item) for key, item in value.items()} + return value + + +@lru_cache(maxsize=None) +def _resolve_checkpoint_dir(model_id_or_path: str) -> Path: + candidate = Path(model_id_or_path).expanduser() + if candidate.exists(): + return candidate + + # Try safetensors first (preferred format for weight-free export) + snapshot_dir = snapshot_download( + repo_id=model_id_or_path, + allow_patterns=["*.safetensors", "*.json"], + ignore_patterns=["*.onnx", "*.ot", "*.md", "*.txt", "*.pdf", "*.msgpack", "*.h5", "*.pth"], + resume_download=True, + ) + snapshot_path = Path(snapshot_dir) + + # Check if any weight files were actually downloaded (not just .json config/tokenizer files) + has_weights = ( + bool(list(snapshot_path.glob("*.safetensors"))) or (snapshot_path / "model.safetensors.index.json").exists() + ) + + if not has_weights: + # Model has no safetensors on Hub — fall back to .bin files. + # CheckpointTransformPipeline.apply() will auto-convert them to safetensors on first use. + snapshot_dir = snapshot_download( + repo_id=model_id_or_path, + allow_patterns=["*.bin", "*.json"], + ignore_patterns=[ + "*.onnx", + "*.ot", + "*.md", + "*.txt", + "*.pdf", + "*.msgpack", + "*.h5", + "*.pth", + "flax_model*", + "tf_model*", + ], + resume_download=True, + ) + + return Path(snapshot_dir) + + +def _resolve_checkpoint_files(model_id_or_path: str) -> List[str]: + checkpoint_dir = _resolve_checkpoint_dir(model_id_or_path) + checkpoint_files = sorted(str(path) for path in checkpoint_dir.glob("*.safetensors")) + if not checkpoint_files: + raise FileNotFoundError(f"No safetensors checkpoint files found for {model_id_or_path}") + return checkpoint_files + + +def _module_name_map(model: nn.Module) -> Dict[int, str]: + return {id(module): name for name, module in model.named_modules()} + + +def _collect_tied_weights(model: nn.Module) -> List[TiedWeightAlias]: + if not getattr(model.config, "tie_word_embeddings", False): + return [] + + input_embeddings = model.get_input_embeddings() + output_embeddings = model.get_output_embeddings() + if input_embeddings is None or output_embeddings is None: + return [] + + module_names = _module_name_map(model) + canonical_name = module_names.get(id(input_embeddings)) + alias_name = module_names.get(id(output_embeddings)) + if not canonical_name or not alias_name or canonical_name == alias_name: + return [] + + return [TiedWeightAlias(alias=f"{alias_name}.weight", canonical=f"{canonical_name}.weight")] + + +def _build_meta_qeff_model(qeff_model): + model_ref = qeff_model.hash_params.get("pretrained_model_name_or_path") + if not model_ref: + raise ValueError( + "Weight-free export requires checkpoint metadata. " + "Pass `pretrained_model_name_or_path=...` when constructing the QEff model manually." + ) + + quant_config = getattr(qeff_model.model.config, "quantization_config", None) + + config = copy.deepcopy(qeff_model.model.config) + source_model = getattr(qeff_model.model, "model", qeff_model.model) + with init_empty_weights(): + try: + meta_model = qeff_model._hf_auto_class.from_config(config, attn_implementation="eager") + except ValueError: + meta_model = source_model.__class__(config) + + if qeff_model.__class__.__name__ == "QEffCausalLMForTextImageToTextModel" and hasattr(meta_model, "vision_tower"): + meta_model.vision_tower = nn.Identity() + + if quant_config is None: + target_dtype = getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.float32 + if target_dtype == torch.bfloat16: + target_dtype = torch.float16 + meta_model = meta_model.to(dtype=target_dtype) + + if not hasattr(meta_model, "get_qeff_language_decoder") and not hasattr(meta_model, "get_qeff_vision_encoder"): + from QEfficient.transformers.models.pytorch_transforms import KVCacheExternalModuleMapperTransform + + meta_model, _ = KVCacheExternalModuleMapperTransform.apply(meta_model) + + meta_qeff_model = qeff_model.__class__( + meta_model, + continuous_batching=getattr(qeff_model, "continuous_batching", False), + qaic_config=copy.deepcopy(getattr(qeff_model.model, "qaic_config", None)), + max_seq_len_cached=getattr(qeff_model.model.config, "max_seq_len_cached", None), + pretrained_model_name_or_path=model_ref, + enable_proxy=getattr(qeff_model, "_enable_proxy", False), + ) + meta_qeff_model.hash_params.update(copy.deepcopy(qeff_model.hash_params)) + + if isinstance(qeff_model.model, PooledModel): + meta_qeff_model.model, _ = PoolingTransform.apply(meta_qeff_model.model, qeff_model.model.pooling_fn) + + if quant_config is not None: + # For quantized models the meta model must use the same quantized layer types as the + # checkpoint so that ONNX initializer names match the checkpoint's storage keys. + # We apply the quantizer's architecture preprocessing (layer-type replacement only, + # no weight loading) AFTER __init__ so that Mxfp4GptOssExpertDequantizeTransform — + # which is part of _pytorch_transforms and targets QEffMxfp4GptOssExperts — has + # already run as a no-op and will not undo the replacement below. + from QEfficient.transformers.quantizers.auto import ( + QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING, + QEFF_AUTO_QUANTIZER_MAPPING, + ) + + # quantization_config may be a plain dict (AutoConfig.from_pretrained) or a proper + # config object (QEFFAutoModelForCausalLM.from_pretrained). Normalise to an object. + if isinstance(quant_config, dict): + quant_type = quant_config.get("quant_method") or quant_config.get("quant_type") + config_cls = QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING.get(quant_type) + if config_cls is None: + raise NotImplementedError( + f"Weight-free export is not implemented for quantization type '{quant_type}'. Supported: mxfp4" + ) + init_kwargs = {k: v for k, v in quant_config.items() if k != "quant_method"} + quant_config = config_cls(**init_kwargs) + else: + quant_method = getattr(quant_config, "quant_method", None) or getattr(quant_config, "quant_type", None) + quant_type = quant_method.value if hasattr(quant_method, "value") else quant_method + + quantizer_cls = QEFF_AUTO_QUANTIZER_MAPPING.get(quant_type) if quant_type else None + if quantizer_cls is None: + raise NotImplementedError( + f"Weight-free export is not implemented for quantization type '{quant_type}'. Supported: mxfp4" + ) + quantizer = quantizer_cls(quant_config) + # Run inside init_empty_weights so newly created quantized layer buffers stay on + # the meta device and are treated as weight-spec entries, not embedded constants. + with init_empty_weights(): + quantizer._process_model_before_weight_loading(meta_qeff_model.model) + + meta_qeff_model.model.eval() + return meta_qeff_model + + +def _checkpoint_root(model_id_or_path: str, checkpoint_files: Sequence[str]) -> Optional[Path]: + if not checkpoint_files: + return None + + candidate = Path(model_id_or_path).expanduser() + if candidate.exists(): + return candidate.parent + + first_checkpoint = Path(checkpoint_files[0]) + for parent in first_checkpoint.parents: + if parent.name.startswith("models--"): + return parent.parent + return first_checkpoint.parent + + +def _load_checkpoint_index(checkpoint_files: List[str]) -> Dict[str, str]: + tensor_to_file = {} + for checkpoint_file in checkpoint_files: + handle = safe_open(checkpoint_file, framework="pt") + for key in handle.keys(): + tensor_to_file[key] = checkpoint_file + return tensor_to_file + + +def _build_location( + checkpoint_files: Sequence[str], + checkpoint_file: Optional[str], + tensor_key: str, +) -> Optional[WeightSpecLocation]: + if checkpoint_file is None: + return None + + return WeightSpecLocation(file=list(checkpoint_files).index(checkpoint_file), key=tensor_key) + + +def _find_checkpoint_key( + onnx_name: str, + checkpoint_index: Dict[str, str], + backbone: nn.Module, +) -> Optional[str]: + """ + Resolve an ONNX initializer name to its key in the safetensors checkpoint. + + Four lookups are attempted in order: + + 1. Direct match — ONNX name == checkpoint key. + Covers decoder-only LLMs (Llama, GPT-OSS) and any model exported + without extra wrappers. + + 2. Strip our PooledModel prefix ("base_model.") and try the bare key. + Covers embedding models whose checkpoint was saved from the base class + directly (e.g. BAAI/bge-base saved from BertModel → bare keys). + + 3. Strip "base_model." then prepend backbone.base_model_prefix. + Covers embedding models whose checkpoint was saved from a task-specific + class (e.g. BAAI/bge-reranker saved from XLMRobertaForSequenceClassification + → "roberta." prefix). HuggingFace defines base_model_prefix on every + model class as exactly the attribute name the task class uses to store + the backbone, so this is generalized — no per-model-family list needed. + + 4. Strip backbone.base_model_prefix from the front of the ONNX name. + Covers ForCausalLM models (e.g. GPT2LMHeadModel) whose checkpoint was + saved from the base model class (e.g. GPT2Model) so keys lack the LM + wrapper prefix — ONNX name "transformer.wte.weight" maps to checkpoint + key "wte.weight" after stripping the "transformer." prefix. + """ + # 1. Direct match + if onnx_name in checkpoint_index: + return onnx_name + + # Strip PooledModel's "base_model." wrapper prefix + stripped = onnx_name[len("base_model.") :] if onnx_name.startswith("base_model.") else onnx_name + + # 2. Bare key — checkpoint saved from base class + if stripped in checkpoint_index: + return stripped + + if stripped.startswith("model."): + without_wrapper = stripped[len("model.") :] + if without_wrapper in checkpoint_index: + return without_wrapper + + prefix = getattr(backbone, "base_model_prefix", "") + + # 3. Task-class prefix — checkpoint saved from ForSequenceClassification etc. + if prefix: + prefixed = f"{prefix}.{stripped}" + if prefixed in checkpoint_index: + return prefixed + + # 4. Strip base_model_prefix from ONNX name — checkpoint saved from the + # base model class so keys lack the ForCausalLM wrapper prefix + # (e.g. GPT2: ONNX "transformer.wte.weight" → checkpoint "wte.weight"). + if prefix and stripped.startswith(f"{prefix}."): + without_prefix = stripped[len(f"{prefix}.") :] + if without_prefix in checkpoint_index: + return without_prefix + + # 5. MLP attribute rename: transformers 5.x renamed 'block_sparse_moe' → 'mlp' + # for Mixtral. The original checkpoint uses 'block_sparse_moe' but the model + # (and ONNX) uses 'mlp'. Try the reverse substitution. + if ".mlp." in stripped: + candidate = stripped.replace(".mlp.", ".block_sparse_moe.") + if candidate in checkpoint_index: + return candidate + + # 6. MoE router attribute rename: QEff wrappers may call the router 'gate' + # while the checkpoint stores it as 'router' (e.g. Qwen3-MoE). + if stripped.endswith(".mlp.gate.weight"): + candidate = stripped[: -len(".gate.weight")] + ".router.weight" + if candidate in checkpoint_index: + return candidate + # Reverse: checkpoint calls it 'gate', ONNX uses 'router' + if stripped.endswith(".mlp.router.weight"): + candidate = stripped[: -len(".router.weight")] + ".gate.weight" + if candidate in checkpoint_index: + return candidate + + return None + + +def _promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name: str, qeff_model) -> WeightSpec: + from QEfficient.transformers.embeddings.embedding_utils import PooledModel + + model_ir = onnx_program.model + model_names = {name for name, _ in qeff_model.model.named_parameters()} + model_names.update({name for name, _ in qeff_model.model.named_buffers()}) + tied_weights = _collect_tied_weights(qeff_model.model) + tied_weight_map = {entry.alias: entry.canonical for entry in tied_weights} + checkpoint_files = _resolve_checkpoint_files(model_ref) + checkpoint_root = _checkpoint_root(model_ref, checkpoint_files) + checkpoint_index = _load_checkpoint_index(checkpoint_files) + relative_checkpoint_files = [ + ExternalDataFile( + path=str(Path(checkpoint_file).relative_to(checkpoint_root)) + if checkpoint_root is not None + else Path(checkpoint_file).name, + format="safetensors", + ) + for checkpoint_file in checkpoint_files + ] + backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model + promoted_inputs: List[WeightSpecInput] = [] + + for name, init_value in list(model_ir.graph.initializers.items()): + onnx_name = tied_weight_map.get(name, name) + checkpoint_key = _find_checkpoint_key(onnx_name, checkpoint_index, backbone) + + if checkpoint_key is None and name in model_names: + checkpoint_key = _find_checkpoint_key(name, checkpoint_index, backbone) + + if checkpoint_key is None: + # Computed buffer (e.g. sin_cached, cos_cached) — leave as ONNX initializer. + # The compiler embeds it in the model; it is not loaded from a checkpoint file. + continue + + location = _build_location(checkpoint_files, checkpoint_index[checkpoint_key], checkpoint_key) + + model_ir.graph.inputs.append( + ir.Value( + name=name, + shape=init_value.shape, + type=ir.TensorType(init_value.dtype), + ) + ) + del model_ir.graph.initializers[name] + promoted_inputs.append(WeightSpecInput(name=name, location=location)) + + return WeightSpec( + model_name=model_name, + model_id=model_ref, + files=relative_checkpoint_files, + inputs=promoted_inputs, + ) + + +def _prune_unused_fake_initializers(onnx_program) -> None: + """Remove FakeTensor initializers not referenced by any graph node. + + During weight-free dynamo export, meta-device parameters that are not + actually consumed in the forward graph can end up as orphan initializers. + Serialising them fails (no data to write), so we drop them here before save. + """ + from torch._subclasses.fake_tensor import FakeTensor + + initializers = onnx_program.model.graph.initializers + used_names = {name for node in onnx_program.model.graph for name in node.inputs} + used_names.update(output.name for output in onnx_program.model.graph.outputs) + for name in list(initializers): + const_value = getattr(initializers[name], "const_value", None) + raw_value = getattr(const_value, "raw", None) + if isinstance(raw_value, FakeTensor) and name not in used_names: + del initializers[name] + + +def _upsert_metadata_prop(model, key: str, value: str) -> None: + """Insert or update a metadata_props entry on an ONNX model. + + Used to embed weight_spec.json into the ONNX so the QAIC compiler + can locate external weight files without a separate sidecar lookup. + """ + import onnx + + for entry in model.metadata_props: + if entry.key == key: + entry.value = value + return + model.metadata_props.append(onnx.StringStringEntryProto(key=key, value=value)) + + +def _set_module_tensor(module: nn.Module, tensor_name: str, tensor: torch.Tensor) -> None: + parts = tensor_name.split(".") + parent = module + for part in parts[:-1]: + parent = parent[int(part)] if part.isdigit() else getattr(parent, part) + leaf_name = parts[-1] + existing = getattr(parent, leaf_name) + if isinstance(existing, torch.nn.Parameter): + setattr(parent, leaf_name, torch.nn.Parameter(tensor, requires_grad=existing.requires_grad)) + else: + setattr(parent, leaf_name, tensor) + + +def _materialize_checkpoint_metadata_tensors(qeff_model, model_ref: str) -> None: + checkpoint_files = _resolve_checkpoint_files(model_ref) + checkpoint_index = _load_checkpoint_index(checkpoint_files) + backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model + tensors = list(qeff_model.model.named_parameters()) + list(qeff_model.model.named_buffers()) + for name, tensor in tensors: + if not name.endswith("weight_shape") or not getattr(tensor, "is_meta", False): + continue + checkpoint_key = _find_checkpoint_key(name, checkpoint_index, backbone) + if checkpoint_key is None: + continue + checkpoint_file = checkpoint_index[checkpoint_key] + with safe_open(checkpoint_file, framework="pt", device="cpu") as handle: + concrete_tensor = handle.get_tensor(checkpoint_key) + _set_module_tensor(qeff_model.model, name, concrete_tensor) + + +def export_loaded_model_weight_free_onnx( + qeff_model, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_shapes: Dict[str, Any], + export_kwargs: Dict[str, Any], + onnx_transform_kwargs: Dict[str, Any], +): + """Export an already-transformed model, then externalize its initializers. + + This Kimi fallback avoids compressed-tensors meta tracing limitations while + still producing a weight-free ONNX plus weight_spec.json. + """ + onnx_program = torch.onnx.export( + qeff_model.model, + args=(), + f=None, + kwargs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=None, + dynamic_shapes=dynamic_shapes, + **export_kwargs, + ) + if onnx_program is None: + raise RuntimeError("torch.onnx.export returned None for loaded-model weight-free export") + + checkpoint_dir = tmp_onnx_path.parent / "loaded_weight_checkpoint" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / "model.safetensors" + state_dict = { + name: tensor.detach().cpu().contiguous().clone() + for name, tensor in qeff_model.model.state_dict().items() + if torch.is_tensor(tensor) and not tensor.is_meta + } + save_file(state_dict, str(checkpoint_path)) + + spec = _promote_initializers_and_build_spec( + onnx_program, + model_ref=str(checkpoint_dir), + model_name=qeff_model.model_name, + qeff_model=qeff_model, + ) + save_weight_spec(tmp_onnx_path.with_name("weight_spec.json"), spec) + onnx_program.save(str(tmp_onnx_path)) + + def _cleanup(): + return None + + return qeff_model, onnx_transform_kwargs, _cleanup + + +def export_weight_free_onnx( + qeff_model, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_shapes: Dict[str, Any], + export_kwargs: Dict[str, Any], + onnx_transform_kwargs: Dict[str, Any], +): + meta_qeff_model = _build_meta_qeff_model(qeff_model) + cleanup_required = False + + if getattr(qeff_model, "_use_onnx_subfunctions", False): + _, subfunc_kwargs, _ = _setup_onnx_subfunctions( + meta_qeff_model, + (), + { + "use_dynamo": True, + "onnx_transform_kwargs": copy.deepcopy(onnx_transform_kwargs), + "output_names": list(output_names), + }, + ) + onnx_transform_kwargs = subfunc_kwargs.get("onnx_transform_kwargs", onnx_transform_kwargs) + cleanup_required = True + + decoder_layer_classes = get_decoder_layer_classes_for_export(meta_qeff_model.model) + if getattr(meta_qeff_model, "_use_onnx_subfunctions", False) and decoder_layer_classes: + export_context = temporarily_enable_nested_compile_regions(meta_qeff_model.model, decoder_layer_classes) + else: + export_context = temporarily_disable_nested_compile_regions(meta_qeff_model.model, decoder_layer_classes) + + meta_example_inputs = _to_meta(example_inputs) + model_ref = meta_qeff_model.hash_params["pretrained_model_name_or_path"] + _materialize_checkpoint_metadata_tensors(meta_qeff_model, model_ref) + + meta_qeff_model.model.requires_grad_(False) + with export_context: + onnx_program = torch.onnx.export( + meta_qeff_model.model, + args=(), + f=None, + kwargs=meta_example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=None, + dynamic_shapes=dynamic_shapes, + **export_kwargs, + ) + if onnx_program is None: + raise RuntimeError("torch.onnx.export returned None for weight-free dynamo export") + + # Prepare checkpoint: stack MoE experts (if needed) and convert dtype. + # Store next to the SOURCE checkpoint directory (not inside the hashed export dir) + # so any model config variant pointing at the same source reuses the prepared data. + # Include the dtype in the directory name to avoid collisions between fp16/fp32 exports. + target_dtype = ( + getattr(qeff_model.model.config, "dtype", None) + or getattr(qeff_model.model.config, "torch_dtype", None) + or torch.float32 + ) + if target_dtype == torch.bfloat16: + target_dtype = torch.float16 + dtype_suffix = str(target_dtype).replace("torch.", "") # "float16" or "float32" + prep_pipeline = CheckpointTransformPipeline(transforms=qeff_model._checkpoint_transforms) + source_dir = _resolve_checkpoint_dir(model_ref) + prepared_out = source_dir.parent / (source_dir.name + f"-qeff-prepared-{dtype_suffix}") + + _prep_profiler = _QEffMemoryProfiler(sampling_interval=0.05, verbose=False) + _prep_profiler.start_monitoring() + _prep_profiler.mark_operation("Checkpoint Prep") + + prepared_model_ref = str( + prep_pipeline.apply( + src=source_dir, + out=prepared_out, + target_dtype=target_dtype, + ) + ) + + _prep_profiler.stop_monitoring() + print(_prep_profiler.get_memory_report()) + logger.info(_prep_profiler.get_memory_report()) + + # Expose prep stats as module-level variables so external runners + # (e.g. compare_weightfree._runner.py) can read them without parsing stdout. + import QEfficient.exporter.weight_free.core as _self_module + + _self_module._last_prep_peak_rss_mb = _prep_profiler.peak_rss + _self_module._last_prep_duration_seconds = _prep_profiler.operation_durations.get("Checkpoint Prep", 0.0) + _self_module._checkpoint_prep_ran = True + + spec = _promote_initializers_and_build_spec( + onnx_program=onnx_program, + model_ref=prepared_model_ref, + model_name=qeff_model.model_name, + qeff_model=meta_qeff_model, + ) + _prune_unused_fake_initializers(onnx_program) + onnx_program.save(str(tmp_onnx_path)) + save_weight_spec(resolve_weight_spec_path(tmp_onnx_path), spec) + + def cleanup(): + if cleanup_required: + _cleanup_onnx_subfunctions(meta_qeff_model) + + return meta_qeff_model, onnx_transform_kwargs, cleanup + + +def _load_checkpoint_tensor(checkpoint_file: str, key: str) -> np.ndarray: + handle = safe_open(checkpoint_file, framework="pt") + tensor = handle.get_tensor(key).detach().cpu() + # numpy does not support bfloat16; cast to float32 for ORT compatibility + if tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) + return tensor.numpy() + + +def _default_weights_roots(weight_spec_path: Path, spec) -> List[Path]: + roots = [] + ext_root = os.environ.get("AIC_EXTERNAL_DATA_ROOT") + if ext_root: + roots.append(Path(ext_root).expanduser()) + roots.append(weight_spec_path.parent) + candidate = Path(spec.model_id).expanduser() + if candidate.exists(): + roots.append(candidate.parent) + else: + checkpoint_dir = _resolve_checkpoint_dir(spec.model_id) + checkpoint_root = _checkpoint_root(spec.model_id, [str(path) for path in checkpoint_dir.glob("*.safetensors")]) + if checkpoint_root is not None: + roots.append(checkpoint_root) + + deduped_roots: List[Path] = [] + seen = set() + for root in roots: + resolved = root.resolve() + if resolved in seen: + continue + seen.add(resolved) + deduped_roots.append(resolved) + return deduped_roots + + +def _resolve_location_file( + location: WeightSpecLocation, + files: Sequence[ExternalDataFile], + candidate_roots: Sequence[Path], +) -> Path: + if isinstance(location.file, int): + location_path = Path(files[location.file].path) + else: + location_path = Path(location.file) + if location_path.is_absolute(): + return location_path + + for root in candidate_roots: + candidate = root / location_path + if candidate.exists(): + return candidate + + return candidate_roots[0] / location_path if candidate_roots else location_path + + +def load_weight_free_ort_inputs( + weight_spec_path: Path, + runtime_inputs: Dict[str, np.ndarray], + weights_root: Optional[Path] = None, +) -> Dict[str, np.ndarray]: + weight_spec_path = Path(weight_spec_path) + spec = load_weight_spec(weight_spec_path) + candidate_roots = [] + if weights_root is not None: + candidate_roots.append(Path(weights_root).expanduser().resolve()) + candidate_roots.extend(_default_weights_roots(weight_spec_path, spec)) + + ort_inputs = dict(runtime_inputs) + for spec_input in spec.inputs: + if spec_input.name in ort_inputs: + continue + checkpoint_file = _resolve_location_file(spec_input.location, spec.files, candidate_roots) + ort_inputs[spec_input.name] = _load_checkpoint_tensor(str(checkpoint_file), spec_input.location.key) + + return ort_inputs + + +def log_weight_free_export(onnx_path: Path) -> None: + logger.info(f"Weight-free ONNX exported to {onnx_path} with spec {resolve_weight_spec_path(onnx_path)}") diff --git a/QEfficient/exporter/weight_free/spec.py b/QEfficient/exporter/weight_free/spec.py new file mode 100644 index 0000000000..995ab328e0 --- /dev/null +++ b/QEfficient/exporter/weight_free/spec.py @@ -0,0 +1,98 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import json +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Union + +WEIGHT_SPEC_VERSION = 5 + + +@dataclass +class TiedWeightAlias: + alias: str + canonical: str + + +@dataclass +class ExternalDataFile: + path: str + format: str + + +CheckpointFile = ExternalDataFile + + +@dataclass +class WeightSpecLocation: + file: Union[int, str] + key: str + + +@dataclass +class WeightSpecInput: + name: str + location: WeightSpecLocation # required: every spec entry must point to a file + + +@dataclass +class WeightSpec: + model_name: str + model_id: str + files: List[ExternalDataFile] = field(default_factory=list) + inputs: List[WeightSpecInput] = field(default_factory=list) + version: int = WEIGHT_SPEC_VERSION + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + data["model_id"] = str(data["model_id"]) + return data + + +def save_weight_spec(path: Path, spec: WeightSpec) -> Path: + with path.open("w", encoding="utf-8") as handle: + json.dump(spec.to_dict(), handle, indent=2, sort_keys=True) + return path + + +def _load_files(raw: list) -> List[ExternalDataFile]: + if not raw: + return [] + # Backward compat: old format stored plain strings + if isinstance(raw[0], str): + return [ExternalDataFile(path=entry, format="safetensors") for entry in raw] + return [ExternalDataFile(**entry) for entry in raw] + + +def _load_location(raw: dict) -> WeightSpecLocation: + # Backward compat: old format had a redundant "type" field on the location + return WeightSpecLocation(file=raw["file"], key=raw["key"]) + + +def load_weight_spec(path: Path) -> WeightSpec: + with path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + + return WeightSpec( + model_name=data["model_name"], + model_id=data["model_id"], + files=_load_files(data.get("files", data.get("checkpoint_files", []))), + inputs=[ + WeightSpecInput( + name=entry["name"], + location=_load_location(entry["location"]), + ) + for entry in data["inputs"] + if entry.get("location") is not None # backward compat: skip old buffer-only entries + ], + version=data.get("version", WEIGHT_SPEC_VERSION), + ) + + +def resolve_weight_spec_path(onnx_path: Path) -> Path: + return onnx_path.with_name("weight_spec.json") diff --git a/QEfficient/exporter/weight_free/transforms.py b/QEfficient/exporter/weight_free/transforms.py new file mode 100644 index 0000000000..46c9e8db62 --- /dev/null +++ b/QEfficient/exporter/weight_free/transforms.py @@ -0,0 +1,905 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +""" +Checkpoint transforms for weight-free ONNX export. + +Two transforms are provided, selected in priority order by CheckpointTransformPipeline: + + MoEExpertStackingCheckpointTransform — stacks per-expert HF keys into batched + tensors AND converts dtype, all in a single read pass over the checkpoint. + Skipped automatically (is_applicable=False) for dense models. + + DtypeConversionCheckpointTransform — converts all floating-point tensors to + target_dtype. Used as the fallback path for dense models. + +Usage on model classes (mirrors _pytorch_transforms / _onnx_transforms):: + + _checkpoint_transforms = [ + MoEExpertStackingCheckpointTransform, # no-op for dense (is_applicable=False) + DtypeConversionCheckpointTransform, # fallback for dense models + ] +""" + +import json +import os +import re +import shutil +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Type + +import psutil +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from QEfficient.transformers.quantizers.quantizer_utils import convert_moe_packed_tensors +from QEfficient.utils.logging_utils import logger + +# --------------------------------------------------------------------------- +# System-state helpers — used to derive worker counts at runtime +# --------------------------------------------------------------------------- + + +def _available_ram_gb() -> float: + """Available (free + reclaimable) RAM on the current machine in GB.""" + return psutil.virtual_memory().available / 1024**3 + + +def _cpu_count() -> int: + """Logical CPU count with a safe fallback.""" + return os.cpu_count() or 8 + + +def _estimate_layer_stack_gb( + expert_entries: Dict[Tuple[int, int, str], Tuple[str, str]], + layer_idx: int, + num_experts: int, + src: Path, + target_dtype: torch.dtype = torch.float32, +) -> float: + """Estimate peak RAM (GB) required to stack one MoE layer's experts. + + At the moment stacker.stack(target_dtype) runs, five tensors exist in RAM: + + Inputs (checkpoint dtype, e.g. BF16): + gate [E, I, H] + up [E, I, H] + down [E, H, I] + + Outputs (target_dtype, e.g. FP32 — twice as large when converting BF16→FP32): + gate_up [E, 2I, H] cat(gate, up).to(target_dtype) + down_out [E, H, I] down.to(target_dtype) + + Using source dtype bytes for the outputs underestimates by ~45% when + converting BF16→FP32, causing too many parallel workers and OOM. + Returns 1.0 GB as a safe fallback if the shape cannot be read. + """ + sample = next( + (v for (li, ei, k), v in expert_entries.items() if li == layer_idx and k in ("gate_proj", "linear", "w1")), + None, + ) + if sample is None: + return 1.0 + + shard_name, orig_key = sample + try: + with safe_open(str(src / shard_name), framework="pt") as f: + sl = f.get_slice(orig_key) + shape = sl.get_shape() # [I, H] + dtype_str = sl.get_dtype() + except Exception: + return 1.0 + + src_bytes = {"F32": 4, "F16": 2, "BF16": 2, "I8": 1}.get(dtype_str, 2) + tgt_bytes = {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(target_dtype, 4) + ffn_dim, hidden_dim = shape + + # Three input accumulators in source dtype + two output tensors in target dtype + input_elements = num_experts * ( + ffn_dim * hidden_dim + ffn_dim * hidden_dim + hidden_dim * ffn_dim + ) # gate + up + down + output_elements = num_experts * (2 * ffn_dim * hidden_dim + hidden_dim * ffn_dim) # gate_up + down_out + return (input_elements * src_bytes + output_elements * tgt_bytes) / 1024**3 + + +# --------------------------------------------------------------------------- +# Auxiliary file names copied alongside the prepared checkpoint +# --------------------------------------------------------------------------- +_AUX_FILES = [ + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "special_tokens_map.json", + "chat_template.jinja", + "vocab.json", + "merges.txt", +] + +_SENTINEL = ".checkpoint_prepared" + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _convert_bin_to_safetensors(src: Path) -> None: + """Load a .bin checkpoint via transformers and re-save as safetensors in place. + + Uses save_pretrained(safe_serialization=True) which correctly handles tied + weights and multi-shard layouts, writing model.safetensors (single file) or + model-NNNNN-of-MMMMM.safetensors + model.safetensors.index.json (multi-shard). + Idempotent — the caller already verified no safetensors files exist before calling. + """ + import gc + + from transformers import AutoConfig, AutoModelForCausalLM + + logger.info(f"No safetensors files found in {src}. Auto-converting .bin → safetensors (one-time).") + config = AutoConfig.from_pretrained(str(src), trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + str(src), + config=config, + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + model.save_pretrained(str(src), safe_serialization=True) + del model + gc.collect() + logger.info(f"Conversion complete — safetensors files written to {src}") + + +def _read_weight_map(src: Path) -> Dict[str, str]: + """Return {tensor_key: shard_filename} from model.safetensors.index.json, + or by scanning all *.safetensors for single-file checkpoints.""" + index_path = src / "model.safetensors.index.json" + if index_path.exists(): + return json.loads(index_path.read_text())["weight_map"] + shard_files = sorted(src.glob("*.safetensors")) + if not shard_files: + raise FileNotFoundError(f"No safetensors files found in {src}") + weight_map: Dict[str, str] = {} + for sf in shard_files: + with safe_open(str(sf), framework="pt") as f: + for k in f.keys(): + weight_map[k] = sf.name + return weight_map + + +def _atomic_save(tensors: Dict[str, torch.Tensor], dst: Path) -> None: + tmp = dst.with_suffix(dst.suffix + ".tmp") + save_file({k: v.contiguous() for k, v in tensors.items()}, str(tmp)) + tmp.replace(dst) + + +def _write_index(out: Path, weight_map: Dict[str, str]) -> None: + files = set(weight_map.values()) + total_size = sum((out / f).stat().st_size for f in files if (out / f).exists()) + index = { + "metadata": {"total_size": total_size}, + "weight_map": dict(sorted(weight_map.items())), + } + (out / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + + +def _copy_aux_files(src: Path, out: Path) -> None: + for name in _AUX_FILES: + src_file = src / name + if src_file.exists() and not (out / name).exists(): + shutil.copy2(str(src_file), str(out / name)) + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +class BaseCheckpointTransform: + """Base class for checkpoint file transforms. Not to be instantiated. + + Each subclass produces a *complete* prepared checkpoint directory in ``out``. + The pipeline picks the first applicable transform and stops — no chaining. + """ + + def __init__(self): + raise TypeError("Checkpoint transform classes are not to be instantiated.") + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> bool: + """Transform checkpoint at ``src``, write result to ``out``. + Returns True if the checkpoint was prepared, False if skipped (idempotent).""" + raise NotImplementedError + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + """Return True if this transform should run for the given checkpoint.""" + return True + + +# --------------------------------------------------------------------------- +# Transform 1: dtype conversion only — dense model path +# --------------------------------------------------------------------------- + + +class DtypeConversionCheckpointTransform(BaseCheckpointTransform): + """Convert all floating-point tensors to ``target_dtype``. + + One pass per shard, shards processed in parallel via ThreadPoolExecutor. + Used as the dense-model fallback; for MoE checkpoints, + MoEExpertStackingCheckpointTransform handles dtype conversion as part + of its own single pass and this transform is never reached. + """ + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("DtypeConversionCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + new_name_for = { + shard: (f"model_{idx:04d}.safetensors" if len(shard_names) > 1 else "model.safetensors") + for idx, shard in enumerate(shard_names) + } + + # I/O-bound: one thread per shard, capped at 4× CPU count and hard-capped + # at 256 — beyond that OS scheduling overhead outweighs I/O parallelism gains. + n_workers = max_workers if max_workers is not None else min(len(shard_names), _cpu_count() * 4, 256) + + def _process_shard(shard_name: str) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_name_for[shard_name]) + + logger.info( + f"DtypeConversionCheckpointTransform: converting {len(shard_names)} shards " + f"→ {target_dtype} | workers={n_workers} (cpus={_cpu_count()})" + ) + with ThreadPoolExecutor(max_workers=n_workers) as ex: + futures = [ex.submit(_process_shard, s) for s in shard_names] + for fut in as_completed(futures): + fut.result() + + new_weight_map = {k: new_name_for[v] for k, v in weight_map.items()} + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"DtypeConversionCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Internal stacker helper for MoE layers +# --------------------------------------------------------------------------- + + +class _LayerStacker: + """Accumulates per-expert tensors for one MoE layer and produces batched output.""" + + def __init__(self, prefix: str, num_experts: int): + self.prefix = prefix + self.num_experts = num_experts + self._gate: Optional[torch.Tensor] = None + self._up: Optional[torch.Tensor] = None + self._down: Optional[torch.Tensor] = None + + def add(self, expert_idx: int, kind: str, tensor: torch.Tensor) -> None: + # Accept qwen3-moe names (gate_proj/up_proj/down_proj), + # grok-1 names (linear/linear_v/linear_1), + # and Mixtral names (w1=gate, w3=up, w2=down) — map to the same accumulators. + if kind in ("gate_proj", "linear", "w1"): + ffn_dim, hidden_dim = tensor.shape + if self._gate is None: + self._gate = torch.empty(self.num_experts, ffn_dim, hidden_dim, dtype=tensor.dtype) + self._gate[expert_idx] = tensor + elif kind in ("up_proj", "linear_v", "w3"): + ffn_dim, hidden_dim = tensor.shape + if self._up is None: + self._up = torch.empty(self.num_experts, ffn_dim, hidden_dim, dtype=tensor.dtype) + self._up[expert_idx] = tensor + else: # down_proj / linear_1 / w2 — shape is [hidden_dim, ffn_dim] + hidden_dim, ffn_dim = tensor.shape + if self._down is None: + self._down = torch.empty(self.num_experts, hidden_dim, ffn_dim, dtype=tensor.dtype) + self._down[expert_idx] = tensor + + def stack(self, target_dtype: torch.dtype) -> Dict[str, torch.Tensor]: + # Output in the exact layout that model __qeff_init__ creates so + # _promote_initializers_and_build_spec finds an exact checkpoint key match. + # _gate [E, I, H] → transpose(1,2) → gate_proj [E, H, I] + # _up [E, I, H] → transpose(1,2) → up_proj [E, H, I] + # _down [E, H, I] → transpose(1,2) → down_proj_t [E, I, H] + gate_proj = self._gate.to(target_dtype).transpose(1, 2).contiguous() + up_proj = self._up.to(target_dtype).transpose(1, 2).contiguous() + down_proj_t = self._down.to(target_dtype).transpose(1, 2).contiguous() + return { + f"{self.prefix}.gate_proj": gate_proj, # [E, H, I] + f"{self.prefix}.up_proj": up_proj, # [E, H, I] + f"{self.prefix}.down_proj_t": down_proj_t, # [E, I, H] + } + + +# --------------------------------------------------------------------------- +# Transform 2: MoE expert stacking + dtype conversion — single pass +# --------------------------------------------------------------------------- + + +class MoEExpertStackingCheckpointTransform(BaseCheckpointTransform): + """Stack per-expert checkpoint keys into batched tensors AND convert dtype. + + Detects the HuggingFace per-expert layout:: + + *.experts.{E}.gate_proj.weight [I, H] x num_experts + *.experts.{E}.up_proj.weight [I, H] x num_experts + *.experts.{E}.down_proj.weight [H, I] x num_experts + + and produces:: + + *.experts.gate_proj [E, H, I] (gate weights, transposed) + *.experts.up_proj [E, H, I] (up weights, transposed) + *.experts.down_proj_t [E, I, H] (down weights, transposed) + + matching the derived parameter layout that QEff MoE model __qeff_init__ + creates, so _promote_initializers_and_build_spec finds an exact key match. + Non-expert keys receive dtype conversion in the same pass. + + Parallelism: + + - Phase 1 (scan): one thread per shard, reads keys only (I/O bound, cheap). + - Phase 2 (stack): one thread per layer, loads and stacks its experts. + - Phase 3 (base): one thread per shard, converts non-expert keys. + + Phases 2 and 3 run concurrently once phase 1 completes. + """ + + EXPERT_RE = re.compile( + r"^(.+\.layers\.(\d+)\..+?\.experts)\.(\d+)\.(gate_proj|up_proj|down_proj|linear|linear_v|linear_1|w1|w2|w3)\.weight$" + ) + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls.EXPERT_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers_scan: Optional[int] = None, + max_workers_layers: Optional[int] = None, + max_workers_base: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("MoEExpertStackingCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + + # ── Phase 1: parallel key scan — no tensor data loaded ──────────────── + # + # expert_entries[(layer_idx, expert_idx, kind)] = (shard_name, orig_key) + # layer_prefix[layer_idx] = prefix up to .experts + # base_entries[orig_key] = shard_name + expert_entries: Dict[Tuple[int, int, str], Tuple[str, str]] = {} + layer_prefix: Dict[int, str] = {} + base_entries: Dict[str, str] = {} + + def _scan(shard_name: str) -> Tuple[Dict, Dict, Dict]: + loc_e: Dict[Tuple[int, int, str], Tuple[str, str]] = {} + loc_p: Dict[int, str] = {} + loc_b: Dict[str, str] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + m = cls.EXPERT_RE.match(key) + if m: + loc_e[(int(m.group(2)), int(m.group(3)), m.group(4))] = (shard_name, key) + loc_p[int(m.group(2))] = m.group(1) + else: + loc_b[key] = shard_name + return loc_e, loc_p, loc_b + + # Phase 1: I/O-bound — cap at 4× logical CPUs, no point exceeding shard count. + # Hard cap at 256: beyond that, OS scheduling overhead outweighs I/O gains. + n_workers_scan = ( + max_workers_scan if max_workers_scan is not None else min(len(shard_names), _cpu_count() * 4, 256) + ) + logger.info( + f"MoEExpertStackingCheckpointTransform: scanning {len(shard_names)} shards " + f"(workers={n_workers_scan}, cpus={_cpu_count()}, ram_avail={_available_ram_gb():.1f} GB)..." + ) + with ThreadPoolExecutor(max_workers=n_workers_scan) as ex: + for loc_e, loc_p, loc_b in ex.map(_scan, shard_names): + expert_entries.update(loc_e) + layer_prefix.update(loc_p) + base_entries.update(loc_b) + + experts_per_layer: Dict[int, set] = {} + for layer_idx, expert_idx, _ in expert_entries: + experts_per_layer.setdefault(layer_idx, set()).add(expert_idx) + layer_indices = sorted(experts_per_layer.keys()) + sample_n = len(next(iter(experts_per_layer.values()))) if experts_per_layer else 0 + logger.info(f" {len(layer_indices)} MoE layers × {sample_n} experts each.") + + new_weight_map: Dict[str, str] = {} + + # ── Phase 2: parallel layer stacking ────────────────────────────────── + # Each layer thread loads its own experts (grouped by shard to open each + # shard at most once per layer), stacks, converts dtype, writes atomically. + def _stack_layer(layer_idx: int) -> Tuple[str, List[str]]: + num_exp = len(experts_per_layer[layer_idx]) + stacker = _LayerStacker(layer_prefix[layer_idx], num_exp) + + # Detect which kind names are present (qwen3-moe: gate_proj/up_proj/down_proj; + # grok-1: linear/linear_v/linear_1). + kinds_present = {k for (li, _, k) in expert_entries if li == layer_idx} + + by_shard: Dict[str, List[Tuple[int, str, str]]] = {} + for exp_idx in range(num_exp): + for kind in kinds_present: + shard_name, orig_key = expert_entries[(layer_idx, exp_idx, kind)] + by_shard.setdefault(shard_name, []).append((exp_idx, kind, orig_key)) + + for shard_name, entries in by_shard.items(): + with safe_open(str(src / shard_name), framework="pt") as f: + for exp_idx, kind, orig_key in entries: + stacker.add(exp_idx, kind, f.get_tensor(orig_key)) + + stacked = stacker.stack(target_dtype) + out_name = f"experts-layer-{layer_idx:05d}.safetensors" + _atomic_save(stacked, out / out_name) + return out_name, list(stacked.keys()) + + # Phase 2: memory-bound — each layer holds all E×3 expert tensors + the + # stacked output in RAM simultaneously. Derive the worker count from + # available RAM so we never OOM: keep 20% headroom, compute RAM per layer + # from the actual tensor shapes in the checkpoint. + if max_workers_layers is not None: + n_workers_layers = max_workers_layers + elif layer_indices: + sample_layer = layer_indices[0] + layer_gb = _estimate_layer_stack_gb( + expert_entries, sample_layer, len(experts_per_layer[sample_layer]), src, target_dtype + ) + available_gb = _available_ram_gb() + usable_gb = available_gb * 0.8 + n_workers_layers = max(1, min(len(layer_indices), int(usable_gb / layer_gb))) + else: + n_workers_layers = 1 + layer_gb = 0.0 + + logger.info( + f" Stacking {len(layer_indices)} layers → {target_dtype} | " + f"workers={n_workers_layers} (~{layer_gb:.2f} GB/layer, " + f"{_available_ram_gb():.1f} GB available)..." + ) + with ThreadPoolExecutor(max_workers=n_workers_layers) as ex: + futures = {ex.submit(_stack_layer, li): li for li in layer_indices} + for fut in as_completed(futures): + li = futures[fut] + out_name, out_keys = fut.result() + for key in out_keys: + new_weight_map[key] = out_name + logger.info(f" layer {li:5d} → {out_name}") + + # ── Phase 3: parallel base shard conversion ──────────────────────────── + by_shard_base: Dict[str, List[str]] = {} + for key, shard_name in base_entries.items(): + by_shard_base.setdefault(shard_name, []).append(key) + + base_shard_list = sorted(by_shard_base) + new_base_name_for = {shard: f"base-{idx:04d}.safetensors" for idx, shard in enumerate(base_shard_list)} + + def _convert_base(shard_name: str, keys: List[str]) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in keys: + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_base_name_for[shard_name]) + + # Phase 3: mixed I/O + memory — one thread per shard, capped at CPU count. + n_workers_base = max_workers_base if max_workers_base is not None else min(len(base_shard_list), _cpu_count()) + logger.info(f" Converting {len(base_shard_list)} base shards → {target_dtype} | workers={n_workers_base}...") + with ThreadPoolExecutor(max_workers=n_workers_base) as ex: + futures_base = [ex.submit(_convert_base, s, keys) for s, keys in by_shard_base.items()] + for fut in as_completed(futures_base): + fut.result() + + for key, shard_name in base_entries.items(): + new_weight_map[key] = new_base_name_for[shard_name] + + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"MoEExpertStackingCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Transform 3: GptOss MXFP4 dequantize + split fused projections +# --------------------------------------------------------------------------- + + +class GptOssMxfp4ExpertDequantSplitCheckpointTransform(BaseCheckpointTransform): + """Dequantize MXFP4-packed stacked expert tensors and split fused gate_up_proj. + + Detects the GptOss MXFP4 checkpoint layout:: + + *.experts.gate_up_proj_blocks [E, 2*I, G, B] U8 + *.experts.gate_up_proj_scales [E, 2*I, G] U8 + *.experts.gate_up_proj_bias [E, 2*I] BF16 + *.experts.down_proj_blocks [E, I, G, B] U8 + *.experts.down_proj_scales [E, I, G] U8 + *.experts.down_proj_bias [E, H] BF16 + + and produces:: + + *.experts.gate_proj [E, H, I] (dequant gate_up_proj, first half) + *.experts.up_proj [E, H, I] (dequant gate_up_proj, second half) + *.experts.gate_proj_bias [E, I] (gate_up_proj_bias, first half) + *.experts.up_proj_bias [E, I] (gate_up_proj_bias, second half) + *.experts.down_proj [E, H, I] (dequant down_proj) + *.experts.down_proj_bias [E, H] (dtype-converted, unchanged key) + + matching the derived parameter layout that QEffGptOssExperts.__qeff_init__ + creates, so _promote_initializers_and_build_spec finds an exact key match. + Non-expert keys receive dtype conversion in the same pass. + + Parallelism mirrors MoEExpertStackingCheckpointTransform: + - Phase 1 (scan): one thread per shard — collect expert tensor locations. + - Phase 2 (dequant): one thread per layer — dequant, split, write. + - Phase 3 (base): one thread per shard — dtype-convert non-expert keys. + """ + + _BLOCKS_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_blocks$") + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls._BLOCKS_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers_scan: Optional[int] = None, + max_workers_layers: Optional[int] = None, + max_workers_base: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("GptOssMxfp4ExpertDequantSplitCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + + # ── Phase 1: scan — collect expert tensor locations ────────────────── + # expert_locs[(layer_idx, kind)] = (blocks_shard, blocks_key, scales_shard, scales_key) + # bias_locs[(layer_idx, kind)] = (shard, key) for gate_up_proj_bias / down_proj_bias + # layer_prefix[layer_idx] = prefix up to .experts + # base_entries[orig_key] = shard_name + _SCALES_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_scales$") + _BIAS_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_bias$") + + expert_locs: Dict[Tuple[int, str], Dict] = {} # {(layer, kind): {blocks/scales: (shard, key)}} + bias_locs: Dict[Tuple[int, str], Tuple[str, str]] = {} + layer_prefix: Dict[int, str] = {} + base_entries: Dict[str, str] = {} + + def _scan(shard_name: str): + loc_e: Dict[Tuple[int, str], Dict] = {} + loc_b: Dict[Tuple[int, str], Tuple[str, str]] = {} + loc_p: Dict[int, str] = {} + loc_base: Dict[str, str] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + m = cls._BLOCKS_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_e.setdefault((li, kind), {})["blocks"] = (shard_name, key) + loc_p[li] = m.group(1) + continue + m = _SCALES_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_e.setdefault((li, kind), {})["scales"] = (shard_name, key) + loc_p[li] = m.group(1) + continue + m = _BIAS_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_b[(li, kind)] = (shard_name, key) + loc_p[li] = m.group(1) + continue + loc_base[key] = shard_name + return loc_e, loc_b, loc_p, loc_base + + n_scan = max_workers_scan if max_workers_scan is not None else min(len(shard_names), _cpu_count() * 4, 256) + logger.info( + f"GptOssMxfp4ExpertDequantSplitCheckpointTransform: scanning {len(shard_names)} shards " + f"(workers={n_scan})..." + ) + with ThreadPoolExecutor(max_workers=n_scan) as ex: + for loc_e, loc_b, loc_p, loc_base in ex.map(_scan, shard_names): + for k, v in loc_e.items(): + expert_locs.setdefault(k, {}).update(v) + bias_locs.update(loc_b) + layer_prefix.update(loc_p) + base_entries.update(loc_base) + + layer_indices = sorted({li for li, _ in expert_locs}) + logger.info(f" Found {len(layer_indices)} MoE layers.") + + new_weight_map: Dict[str, str] = {} + + # ── Phase 2: per-layer dequant + split ──────────────────────────────── + def _process_layer(layer_idx: int) -> Tuple[str, List[str]]: + prefix = layer_prefix[layer_idx] + tensors: Dict[str, torch.Tensor] = {} + + def _load(shard: str, key: str) -> torch.Tensor: + with safe_open(str(src / shard), framework="pt") as f: + return f.get_tensor(key) + + # gate_up_proj: dequant → [E, H, 2*I], then split interleaved (gate=even cols, up=odd cols) + # HF _apply_gate uses gate_up[..., ::2] for gate and gate_up[..., 1::2] for up, + # so columns are interleaved: col 0=gate0, col 1=up0, col 2=gate1, col 3=up1, ... + gu_blocks_shard, gu_blocks_key = expert_locs[(layer_idx, "gate_up_proj")]["blocks"] + gu_scales_shard, gu_scales_key = expert_locs[(layer_idx, "gate_up_proj")]["scales"] + gu_blocks = _load(gu_blocks_shard, gu_blocks_key) + gu_scales = _load(gu_scales_shard, gu_scales_key) + gate_up = convert_moe_packed_tensors(gu_blocks, gu_scales, dtype=target_dtype) + tensors[f"{prefix}.gate_proj"] = gate_up[..., 0::2].contiguous() + tensors[f"{prefix}.up_proj"] = gate_up[..., 1::2].contiguous() + + # gate_up_proj_bias: split [E, 2*I] → [E, I] + [E, I] (same interleaved convention) + if (layer_idx, "gate_up_proj") in bias_locs: + bias_shard, bias_key = bias_locs[(layer_idx, "gate_up_proj")] + gu_bias = _load(bias_shard, bias_key).to(target_dtype) + tensors[f"{prefix}.gate_proj_bias"] = gu_bias[..., 0::2].contiguous() + tensors[f"{prefix}.up_proj_bias"] = gu_bias[..., 1::2].contiguous() + + # down_proj: dequant → [E, H, I] + dp_blocks_shard, dp_blocks_key = expert_locs[(layer_idx, "down_proj")]["blocks"] + dp_scales_shard, dp_scales_key = expert_locs[(layer_idx, "down_proj")]["scales"] + dp_blocks = _load(dp_blocks_shard, dp_blocks_key) + dp_scales = _load(dp_scales_shard, dp_scales_key) + tensors[f"{prefix}.down_proj"] = convert_moe_packed_tensors(dp_blocks, dp_scales, dtype=target_dtype) + + # down_proj_bias: pass through with dtype conversion + if (layer_idx, "down_proj") in bias_locs: + dp_bias_shard, dp_bias_key = bias_locs[(layer_idx, "down_proj")] + tensors[f"{prefix}.down_proj_bias"] = _load(dp_bias_shard, dp_bias_key).to(target_dtype) + + out_name = f"experts-layer-{layer_idx:05d}.safetensors" + _atomic_save(tensors, out / out_name) + return out_name, list(tensors.keys()) + + n_layers = ( + max_workers_layers if max_workers_layers is not None else max(1, min(len(layer_indices), _cpu_count())) + ) + logger.info(f" Dequantizing {len(layer_indices)} layers | workers={n_layers}...") + with ThreadPoolExecutor(max_workers=n_layers) as ex: + futures = {ex.submit(_process_layer, li): li for li in layer_indices} + for fut in as_completed(futures): + li = futures[fut] + out_name, out_keys = fut.result() + for key in out_keys: + new_weight_map[key] = out_name + logger.info(f" layer {li:5d} → {out_name}") + + # ── Phase 3: base shard dtype conversion ────────────────────────────── + by_shard_base: Dict[str, List[str]] = {} + for key, shard_name in base_entries.items(): + by_shard_base.setdefault(shard_name, []).append(key) + + base_shard_list = sorted(by_shard_base) + new_base_name_for = {shard: f"base-{idx:04d}.safetensors" for idx, shard in enumerate(base_shard_list)} + + def _convert_base(shard_name: str, keys: List[str]) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in keys: + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_base_name_for[shard_name]) + + n_base = max_workers_base if max_workers_base is not None else min(len(base_shard_list), _cpu_count()) + logger.info(f" Converting {len(base_shard_list)} base shards | workers={n_base}...") + with ThreadPoolExecutor(max_workers=n_base) as ex: + futures_base = [ex.submit(_convert_base, s, keys) for s, keys in by_shard_base.items()] + for fut in as_completed(futures_base): + fut.result() + + for key, shard_name in base_entries.items(): + new_weight_map[key] = new_base_name_for[shard_name] + + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"GptOssMxfp4ExpertDequantSplitCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + + +class MoEFusedExpertSplitCheckpointTransform(BaseCheckpointTransform): + """Split already-stacked MoE expert weights into the derived layout. + + Some MoE checkpoints (e.g. Mixtral transformers >= 5.x) store experts + as per-layer fused tensors rather than per-expert individual weights: + + *.experts.gate_up_proj [E, 2*I, H] (gate and up concatenated) + *.experts.down_proj [E, H, I] + + The QEff model wrappers create derived parameters that the ONNX + initializer names refer to: + + *.experts.gate_proj [E, H, I] = gate_up_proj[:, :ffn_dim, :].T(1,2) + *.experts.up_proj [E, H, I] = gate_up_proj[:, ffn_dim:, :].T(1,2) + *.experts.down_proj_t [E, I, H] = down_proj.T(1,2) + + is_applicable returns True only when the fused format is detected. + Old-format checkpoints with per-expert keys (e.g. experts.0.gate_proj.weight) + are handled by MoEExpertStackingCheckpointTransform instead. + Also handles dtype conversion in the same pass. + """ + + _FUSED_GATE_UP_RE = re.compile(r"^(.+\.experts)\.gate_up_proj$") + _FUSED_DOWN_RE = re.compile(r"^(.+\.experts)\.down_proj$") + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls._FUSED_GATE_UP_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> bool: + index_path = src / "model.safetensors.index.json" + if index_path.exists(): + weight_map: Dict[str, str] = json.loads(index_path.read_text())["weight_map"] + else: + shards = sorted(src.glob("*.safetensors")) + if not shards: + return False + weight_map = {} + for shard in shards: + with safe_open(str(shard), framework="pt") as f: + for k in f.keys(): + weight_map[k] = shard.name + + if not cls.is_applicable(weight_map): + return False + + out.mkdir(parents=True, exist_ok=True) + + new_weight_map: Dict[str, str] = {} + for shard_name in sorted(set(weight_map.values())): + shard_src = src / shard_name + if not shard_src.exists(): + continue + + out_tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(shard_src), framework="pt") as f: + for key in f.keys(): + tensor = f.get_tensor(key).to(target_dtype) + gate_up_m = cls._FUSED_GATE_UP_RE.match(key) + down_m = cls._FUSED_DOWN_RE.match(key) + + if gate_up_m: + prefix = gate_up_m.group(1) + ffn_dim = tensor.shape[1] // 2 + # Split fused [E,2I,H] → gate/up each [E,H,I] + out_tensors[f"{prefix}.gate_proj"] = tensor[:, :ffn_dim, :].transpose(1, 2).contiguous() + out_tensors[f"{prefix}.up_proj"] = tensor[:, ffn_dim:, :].transpose(1, 2).contiguous() + new_weight_map[f"{prefix}.gate_proj"] = shard_name + new_weight_map[f"{prefix}.up_proj"] = shard_name + # Keep original for completeness + out_tensors[key] = tensor + new_weight_map[key] = shard_name + elif down_m: + prefix = down_m.group(1) + # Transpose [E,H,I] → [E,I,H] + out_tensors[f"{prefix}.down_proj_t"] = tensor.transpose(1, 2).contiguous() + new_weight_map[f"{prefix}.down_proj_t"] = shard_name + out_tensors[key] = tensor + new_weight_map[key] = shard_name + else: + out_tensors[key] = tensor + new_weight_map[key] = shard_name + + save_file({k: v.contiguous() for k, v in out_tensors.items()}, str(out / shard_name)) + + (out / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": new_weight_map}, indent=2) + ) + return True + + +class CheckpointTransformPipeline: + """Selects and runs the first applicable checkpoint transform. + + Transforms are priority-ordered. The first one whose ``is_applicable()`` + returns True is executed and the pipeline stops. Each transform produces a + complete prepared checkpoint — there is no chaining between transforms. + + Example:: + + pipeline = CheckpointTransformPipeline([ + MoEExpertStackingCheckpointTransform, # MoE models: stacks + converts + DtypeConversionCheckpointTransform, # dense models: converts only + ]) + prepared_dir = pipeline.apply(src, out, target_dtype=torch.float32) + """ + + def __init__(self, transforms: List[Type[BaseCheckpointTransform]]): + self.transforms = transforms + + def apply( + self, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> Path: + src, out = Path(src), Path(out) + + # Auto-convert .bin checkpoints to safetensors on first use. + # Idempotent: skipped on subsequent runs once safetensors files exist. + has_safetensors = bool(list(src.glob("*.safetensors"))) or (src / "model.safetensors.index.json").exists() + if not has_safetensors and list(src.glob("*.bin")): + _convert_bin_to_safetensors(src) + + weight_map = _read_weight_map(src) + for transform in self.transforms: + if transform.is_applicable(weight_map): + transform.apply(src, out, target_dtype=target_dtype, **kwargs) + return out + return src # no transform applicable — source is already usable as-is diff --git a/QEfficient/generation/embedding_handler.py b/QEfficient/generation/embedding_handler.py index 2a5a61f6b8..c8f3528094 100755 --- a/QEfficient/generation/embedding_handler.py +++ b/QEfficient/generation/embedding_handler.py @@ -235,18 +235,23 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - image = image.resize( (constants.GRANITEVISION_IMG_SIZE_HEIGHT, constants.GRANITEVISION_IMG_SIZE_WIDTH) ) + model_type = getattr(getattr(self._qeff_model, "model", None).config, "model_type", "") + # Gemma4 expects the processor-rendered prompt with the image placeholder ahead of user text. is_gemma4 = ( hasattr(self._qeff_model.model.config, "model_type") and self._qeff_model.model.config.model_type == "gemma4" ) + is_kimi_k25 = model_type == "kimi_k25" # Prepare conversation format conversation = [ { "role": "user", "content": ( - [{"type": "image"}, {"type": "text", "text": query}] + [{"type": "image_url", "image_url": image}, {"type": "text", "text": query}] + if is_kimi_k25 + else [{"type": "image"}, {"type": "text", "text": query}] if is_gemma4 else [{"type": "text", "text": query}, {"type": "image"}] ), @@ -254,22 +259,29 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - ] # Apply chat template - if is_gemma4: + if is_kimi_k25: + inputs = self._processor( + messages=conversation, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + elif is_gemma4: prompt = self._processor.apply_chat_template( conversation, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) + inputs = self._processor(images=image, text=prompt, return_tensors="pt") else: prompt = self._processor.apply_chat_template( conversation, tokenize=False, add_generation_prompt=True, ) - # Process image and text - inputs = self._processor(images=image, text=prompt, return_tensors="pt") - model_type = getattr(getattr(self._qeff_model, "model", None).config, "model_type", "") + inputs = self._processor(images=image, text=prompt, return_tensors="pt") + if model_type in { "qwen2_5_vl", "qwen3_vl_moe", @@ -299,6 +311,13 @@ def prepare_vlm_inputs(self, image_url: str, query: str, prefill_seq_len: int) - }: vision_inputs[k] = np.array(v) + if is_kimi_k25: + grid_thws = inputs.get("grid_thws") + if grid_thws is None: + raise ValueError("Kimi-K2.5 processor output must include grid_thws for vision export.") + vision_inputs["h_shape"] = np.ones(int(grid_thws[0, 1].item()), dtype=np.int64) + vision_inputs["w_shape"] = np.ones(int(grid_thws[0, 2].item()), dtype=np.int64) + # Convert specific inputs to float16 vision_inputs_fp16 = {"pixel_values", "image_masks"} for k in vision_inputs_fp16: diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index 7af4bfbefa..4f016b2d49 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -423,8 +423,14 @@ def __init__(self, ckv, k_pe): def update_ckv(self, compressed_kv, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.ckv = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv) + else: + self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) ckv_out = self.ckv ctx_len = ckv_out.shape[-2] @@ -434,14 +440,23 @@ def update_ckv(self, compressed_kv, cache_kwargs): invalid_idx_value = InvalidIndexProvider._get_invalid_idx_value() ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ckv_out = ctx_gather(ckv_out, ctx_indices, ctx_len) - ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(ckv_out, dtype=ckv_out.dtype), ckv_out) + if batch_index is not None: + ckv_out = ctx_gather_cb(ckv_out, batch_index, ctx_indices, ctx_len) + else: + ckv_out = ctx_gather(ckv_out, ctx_indices, ctx_len) + ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out) return ckv_out def update_k_pe(self, k_pe_cache, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.k_pe = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + else: + self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) k_pe_out = self.k_pe ctx_len = k_pe_out.shape[-2] @@ -451,14 +466,18 @@ def update_k_pe(self, k_pe_cache, cache_kwargs): invalid_idx_value = InvalidIndexProvider._get_invalid_idx_value() ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - k_pe_out = ctx_gather(k_pe_out, ctx_indices, ctx_len) - k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(k_pe_out, dtype=k_pe_out.dtype), k_pe_out) + if batch_index is not None: + k_pe_out = ctx_gather_cb(k_pe_out, batch_index, ctx_indices, ctx_len) + else: + k_pe_out = ctx_gather(k_pe_out, ctx_indices, ctx_len) + k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), k_pe_out) return k_pe_out def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): # Gather ckv_out = self.ckv position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = ckv_out.shape ctx_indices = torch.arange(start=start_index, end=end_index, dtype=position_ids.dtype)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1).to(position_ids.dtype) @@ -468,8 +487,11 @@ def read_only_blocked_ckv(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - ckv_out = ctx_gather_blocked_kv(ckv_out, ctx_indices) + if batch_index is not None: + ckv_out = ctx_gather_blocked_kv_cb(ckv_out, batch_index, ctx_indices) + else: + ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) + ckv_out = ctx_gather_blocked_kv(ckv_out, ctx_indices) ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(ckv_out, dtype=ckv_out.dtype), ckv_out) return ckv_out @@ -478,6 +500,7 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): # Gather k_pe_out = self.k_pe position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = k_pe_out.shape ctx_indices = torch.arange(start=start_index, end=end_index, dtype=position_ids.dtype)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1).to(position_ids.dtype) @@ -487,22 +510,37 @@ def read_only_blocked_k_pe(self, start_index, end_index, cache_kwargs): ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) - ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) - k_pe_out = ctx_gather_blocked_kv(k_pe_out, ctx_indices) + if batch_index is not None: + k_pe_out = ctx_gather_blocked_kv_cb(k_pe_out, batch_index, ctx_indices) + else: + ctx_indices = ctx_indices.expand(batch, num_kv_heads, ctx_indices.shape[-1]) + k_pe_out = ctx_gather_blocked_kv(k_pe_out, ctx_indices) k_pe_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(k_pe_out, dtype=k_pe_out.dtype), k_pe_out) return k_pe_out def write_only_k_pe(self, k_pe_cache, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.k_pe = ctx_scatter_cb(self.k_pe, batch_index, scatter_position_ids, k_pe_cache) + else: + self.k_pe = ctx_scatter(self.k_pe, position_ids, k_pe_cache) return self.k_pe def write_only_ckv(self, compressed_kv, cache_kwargs): position_ids = cache_kwargs.get("position_ids") + batch_index = cache_kwargs.get("batch_index", None) # Check and fetch batch index value form the kwargs - self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) + if batch_index is not None: + invalid_scatter_index = torch.iinfo(torch.int32).max + scatter_position_ids = torch.where(position_ids < 0, invalid_scatter_index, position_ids) + self.ckv = ctx_scatter_cb(self.ckv, batch_index, scatter_position_ids, compressed_kv) + else: + self.ckv = ctx_scatter(self.ckv, position_ids, compressed_kv) return self.ckv diff --git a/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py b/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py deleted file mode 100644 index 7f68c3d86e..0000000000 --- a/QEfficient/transformers/models/deepseek_v3/configuration_deepseek.py +++ /dev/null @@ -1,219 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ---------------------------------------------------------------------------- - -from transformers.configuration_utils import PretrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - -DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} - - -class DeepseekV3Config(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek - model according to the specified arguments, defining the model architecture. Instantiating a configuration with the - defaults will yield a similar configuration to that of the DeepSeek-V3. - - Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PretrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 129280): - Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`DeepseekV3Model`] - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 11008): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 1407): - Dimension of the MoE representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_nextn_predict_layers (`int`, *optional*, defaults to 1): - Number of nextn predict layers in the DeepSeekV3 Model. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the Transformer decoder. - n_shared_experts (`int`, *optional*, defaults to None): - Number of shared experts, None means dense model. - n_routed_experts (`int`, *optional*, defaults to None): - Number of routed experts, None means dense model. - routed_scaling_factor (`float`, *optional*, defaults to 1.0): - Scaling factor or routed experts. - topk_method (`str`, *optional*, defaults to `gready`): - Topk method used in routed gate. - n_group (`int`, *optional*, defaults to None): - Number of groups for routed experts. - topk_group (`int`, *optional*, defaults to None): - Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). - num_experts_per_tok (`int`, *optional*, defaults to None): - Number of selected experts, None means dense model. - moe_layer_freq (`int`, *optional*, defaults to 1): - The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. - first_k_dense_replace (`int`, *optional*, defaults to 0): - Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). - \--k dense layers--/ - norm_topk_prob (`bool`, *optional*, defaults to False): - Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to 'softmax'): - Method of computing expert weights. - aux_loss_alpha (`float`, *optional*, defaults to 0.001): - Auxiliary loss weight coefficient. - seq_aux = (`bool`, *optional*, defaults to True): - Whether to compute the auxiliary loss for each individual sample. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details checkout [this - paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to - `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 2048): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this - document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is - necessary to ensure exact reproducibility of the pretraining results. Please refer to [this - issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling - strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is - `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update - `max_position_embeddings` to the expected new maximum. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - - ```python - >>> from transformers import DeepseekV3Model, DeepseekV3Config - - >>> # Initializing a Deepseek-V3 style configuration - >>> configuration = DeepseekV3Config() - - >>> # Accessing the model configuration - >>> configuration = model.config - ```""" - - model_type = "deepseek_v3" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=129280, - hidden_size=7168, - intermediate_size=18432, - moe_intermediate_size=2048, - num_hidden_layers=61, - num_nextn_predict_layers=1, - num_attention_heads=128, - num_key_value_heads=128, - n_shared_experts=1, - n_routed_experts=256, - ep_size=1, - routed_scaling_factor=2.5, - kv_lora_rank=512, - q_lora_rank=1536, - qk_rope_head_dim=64, - v_head_dim=128, - qk_nope_head_dim=128, - topk_method="noaux_tc", - n_group=8, - topk_group=4, - num_experts_per_tok=8, - moe_layer_freq=1, - first_k_dense_replace=3, - norm_topk_prob=True, - scoring_func="sigmoid", - aux_loss_alpha=0.001, - seq_aux=True, - hidden_act="silu", - max_position_embeddings=4096, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=None, - bos_token_id=0, - eos_token_id=1, - pretraining_tp=1, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_nextn_predict_layers = num_nextn_predict_layers - self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.ep_size = ep_size - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.topk_method = topk_method - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.moe_layer_freq = moe_layer_freq - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.scoring_func = scoring_func - self.aux_loss_alpha = aux_loss_alpha - self.seq_aux = seq_aux - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index 8f08a3169b..ed1f94d4ce 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -20,6 +20,13 @@ generic_blocked_attention_interface, generic_blocked_mla_attention_interface, ) +from QEfficient.customop.ctx_scatter_gather import ( + CtxGatherFunc3DGeneralized, + CtxScatterFunc3DGeneralized, + CtxScatterFunc3DInt, +) +from QEfficient.customop.matmulnbits import QMOE, QuantLinearTorchFunction +from QEfficient.customop.quantization_ops import cast_to_uint4, dequantize_linear from QEfficient.customop.rms_norm import CustomRMSNormFunc from QEfficient.customop.utils import select_interface from QEfficient.transformers.cache_utils import QEffDynamicCache, QEffDynamicCompressedKVRopeCache @@ -116,16 +123,6 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) - def forward(self, x, seq_len=None): - # x: [bs, num_attention_heads, seq_len, head_size] - if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: - self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) - - return ( - self.cos_cached[:seq_len].to(dtype=x.dtype), - self.sin_cached[:seq_len].to(dtype=x.dtype), - ) - class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding): def __init__( @@ -185,7 +182,7 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb -def orig_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): +def orig_apply_rotary_pos_emb(q, k, cos, sin): # , position_ids, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: @@ -206,8 +203,6 @@ def orig_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ - cos = cos[position_ids].unsqueeze(unsqueeze_dim) - sin = sin[position_ids].unsqueeze(unsqueeze_dim) b, h, s, d = q.shape q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) @@ -273,6 +268,8 @@ def fused_forward_h_blocking( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -294,8 +291,7 @@ def fused_forward_h_blocking( if compressed_kvs is not None: kva = compressed_kvs.update_ckv(kva, self.layer_idx, cache_kwargs) - cos, sin = self.rotary_emb(kva, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -310,6 +306,7 @@ def fused_forward_h_blocking( q_pe=q_pe, kva=kva, k_pe=k_pe, + batch_index=batch_index, per_head_q_up=self.per_head_q_up, per_head_k_up=self.per_head_k_up, per_head_v_up=self.per_head_v_up, @@ -340,10 +337,13 @@ def fused_forward_kv_blocking( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() + # -- KV compression (write to cache) ---------------------------------- compressed_kv = self.kv_a_proj_with_mqa(hidden_states) compressed_kv = compressed_kv.view(bsz, q_len, -1, self.kv_lora_rank + self.qk_rope_head_dim).transpose(1, 2) @@ -361,8 +361,7 @@ def fused_forward_kv_blocking( if compressed_kvs is not None: compressed_kvs.write_only_ckv(kva, self.layer_idx, cache_kwargs) - cos, sin = self.rotary_emb(hidden_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: compressed_kvs.write_only_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -379,13 +378,11 @@ def fused_forward_kv_blocking( dq_qup_kupT = torch.matmul(q_a_proj_out, qup_kupT) else: dq_qup_kupT = torch.matmul(q_a_proj_out, self.fusedqk) - qkupTrope_nope = torch.cat((dq_qup_kupT, q_pe), dim=-1) - query = qkupTrope_nope + query = torch.cat((dq_qup_kupT, q_pe), dim=-1) # [B, num_heads, q_len, d_abs] else: q_nope = torch.bmm(q_a_proj_out, self.q_up) q_nope = q_nope.view(bsz, q_len, self.num_heads, self.qk_nope_head_dim).transpose(1, 2) - qnope_rope = torch.cat((q_nope, q_pe), dim=-1) - query = qnope_rope + query = torch.cat((q_nope, q_pe), dim=-1) blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -395,6 +392,7 @@ def fused_forward_kv_blocking( per_head_k_up_normal=self.per_head_k_up_normal, per_head_v_up=self.per_head_v_up, attention_mask=attention_mask, + batch_index=batch_index, scaling=self.softmax_scale, layer_idx=self.layer_idx, compressed_kvs=compressed_kvs, @@ -422,6 +420,8 @@ def fused_forward_orig( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -453,8 +453,7 @@ def fused_forward_orig( absorption = False # ---- Rotary ---- - cos, sin = self.rotary_emb(q_pe, seq_len=32 * 1024) # Doesn't need q_pe as head_dim is initialized - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) if compressed_kvs is not None: k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs) @@ -504,7 +503,7 @@ def fused_forward_orig( self.k_up.squeeze(0).view(self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim).permute(1, 0, 2) ) k_nope = torch.matmul(kva_expanded, k_up_per_head) - if k_heads <= 1: + if k_heads == 1: k_pe_expanded = ( k_pe_expanded.unsqueeze(1) .expand(-1, self.num_heads, -1, -1, -1) @@ -540,6 +539,8 @@ def fused_forward( use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, mla_absorption: Optional[Dict[str, bool]] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -556,6 +557,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) elif getattr(blocking_config, "mode", None) == "kv": @@ -571,6 +574,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) else: @@ -586,6 +591,8 @@ def fused_forward( use_cache, cache_position, mla_absorption, + cos_cached, + sin_cached, **kwargs, ) @@ -600,6 +607,8 @@ def forward_full_kv( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -627,15 +636,19 @@ def forward_full_kv( k_nope = kv[:, :, :, : self.qk_nope_head_dim] value_states = kv[:, :, :, self.qk_nope_head_dim :] - cos, sin = self.rotary_emb(value_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) query_states = torch.cat((q_nope, q_pe), -1) k_pe_new = k_pe.expand(-1, self.num_heads, -1, -1) key_states = torch.cat((k_nope, k_pe_new), -1) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "batch_index": batch_index, "position_ids": position_ids} + cache_kwargs = { + "sin": sin_cached, + "cos": cos_cached, + "batch_index": batch_index, + "position_ids": position_ids, + } key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale @@ -666,6 +679,8 @@ def forward_full_kv_h_blocking( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() @@ -687,17 +702,14 @@ def forward_full_kv_h_blocking( kv = ( self.kv_b_proj(self.kv_a_layernorm(kva)) - .view( - bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) # TODO : split this matmul #with k_up and v_up + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) .transpose(1, 2) ) k_nope = kv[:, :, :, : self.qk_nope_head_dim] value_states = kv[:, :, :, self.qk_nope_head_dim :] - cos, sin = self.rotary_emb(value_states, seq_len=32 * 1024) - q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos_cached, sin_cached) query_states = torch.cat((q_nope, q_pe), -1) k_pe_new = k_pe.expand(-1, self.num_heads, -1, -1) @@ -734,6 +746,8 @@ def forward( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + cos_cached: Optional[torch.Tensor] = None, + sin_cached: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) @@ -748,6 +762,8 @@ def forward( output_attentions, use_cache, cache_position, + cos_cached, + sin_cached, **kwargs, ) else: @@ -761,34 +777,180 @@ def forward( output_attentions, use_cache, cache_position, + cos_cached, + sin_cached, **kwargs, ) +EXPERT_BLOCKING_NUM_NSP = int(os.environ.get("EXPERT_BLOCKING_NUM_NSP", "16")) +EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) + + +def _build_matched_idx_from_cumsum(T2Ei: torch.Tensor) -> torch.Tensor: + """Build packed->original token index""" + batch_size, seq_len = T2Ei.shape + int32_max = torch.iinfo(torch.int32).max + int32_max_scalar = torch.tensor(int32_max, dtype=torch.int32, device=T2Ei.device) + token_idx = torch.arange(seq_len, dtype=torch.int32, device=T2Ei.device).unsqueeze(0).expand(batch_size, -1) + valid_prefix = torch.cumsum(T2Ei.to(torch.int32), dim=1) + valid_dest = valid_prefix - 1 + scatter_pos = torch.where(T2Ei, valid_dest, int32_max_scalar) + matched_idx = torch.full_like(token_idx, int32_max) + matched_idx = CtxScatterFunc3DInt.apply( + matched_idx.unsqueeze(-1), + scatter_pos, + token_idx.unsqueeze(-1), + ).squeeze(-1) + return matched_idx + + +class QEffDeepseekMoEGate(nn.Module): + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None) + if self.scoring_func == "sigmoid": + scores = logits.sigmoid() + else: + raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}") + + ### select top-k experts + if self.topk_method == "noaux_tc": + assert not self.training + scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0) + group_scores = torch.einsum( + "abc->ab", scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0] + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] + _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False) + topk_weight = scores.gather(1, topk_idx) + else: + raise NotImplementedError(f"insupportable TopK function for MoE gating: {self.topk_method}") + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = torch.einsum("bi->b", topk_weight).unsqueeze(-1) + 1e-20 + + topk_weight = topk_weight / denominator + topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor + + router_probs = tmp_scores + router_weights = scores + return topk_idx, topk_weight, router_probs, router_weights + + class QEffDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): - self.all_gate_proj = torch.nn.Parameter( - torch.cat( - [exp.gate_proj.compressor.decompress_module(exp.gate_proj).T.unsqueeze(0) for exp in self.experts], - dim=0, - ) + if hasattr(self, "all_gate_qweight") and not hasattr(self, "experts"): + return + + # Get common parameters from first expert + first_expert = self.experts[0] + self._qeff_quantized_experts = hasattr(first_expert.gate_proj, "bits") + if not self._qeff_quantized_experts: + return + + self.bits = first_expert.gate_proj.bits + self.group_size = first_expert.gate_proj.group_size + self.act_fn = first_expert.act_fn + assert first_expert.gate_proj.act_order == first_expert.up_proj.act_order == first_expert.down_proj.act_order, ( + "act_order mismatch" ) - self.all_up_proj = torch.nn.Parameter( - torch.cat( - [exp.up_proj.compressor.decompress_module(exp.up_proj).T.unsqueeze(0) for exp in self.experts], dim=0 - ) + self.act_order = first_expert.gate_proj.act_order + + # Store dimensions for dequantization + self.in_features_gate, self.out_features_gate = ( + first_expert.gate_proj.in_features, + first_expert.gate_proj.out_features, ) - self.all_down_proj = torch.nn.Parameter( - torch.cat( - [exp.down_proj.compressor.decompress_module(exp.down_proj).T.unsqueeze(0) for exp in self.experts], - dim=0, - ) + self.in_features_up, self.out_features_up = first_expert.up_proj.in_features, first_expert.up_proj.out_features + self.in_features_down, self.out_features_down = ( + first_expert.down_proj.in_features, + first_expert.down_proj.out_features, ) - self.act_fn = self.experts[0].act_fn - def moe( + # Stack all parameters along a new dimension (expert dimension) + self.all_gate_qweight = torch.nn.Parameter( + torch.stack([exp.gate_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // 2 + ), + requires_grad=False, + ) + self.all_gate_scales = torch.nn.Parameter( + torch.stack([exp.gate_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // self.group_size + ), + requires_grad=False, + ) + # TODO: Since we know qzeros is always 8 -> Just embed this once into the operator as parameter -> explore this later + self.all_gate_qzeros = torch.nn.Parameter( + torch.stack([exp.gate_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_gate_gidx = torch.nn.Parameter( + torch.stack([exp.gate_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + + self.all_up_qweight = torch.nn.Parameter( + torch.stack([exp.up_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // 2 + ), + requires_grad=False, + ) + self.all_up_scales = torch.nn.Parameter( + torch.stack([exp.up_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // self.group_size + ), + requires_grad=False, + ) + self.all_up_qzeros = torch.nn.Parameter( + torch.stack([exp.up_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_up_gidx = torch.nn.Parameter( + torch.stack([exp.up_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + + self.all_down_qweight = torch.nn.Parameter( + torch.stack([exp.down_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // 2 + ), + requires_grad=False, + ) + self.all_down_scales = torch.nn.Parameter( + torch.stack([exp.down_proj.scales for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // self.group_size + ), + requires_grad=False, + ) + self.all_down_qzeros = torch.nn.Parameter( + torch.stack([exp.down_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_down_gidx = torch.nn.Parameter( + torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + del self.experts + + def moe_old( self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, @@ -798,29 +960,174 @@ def moe( hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - gate_proj = self.all_gate_proj[topk_indices.flatten()] - up_proj = self.all_up_proj[topk_indices.flatten()] - down_proj = self.all_down_proj[topk_indices.flatten()] - expert_in = ( - hidden_states.unsqueeze(1).expand(-1, self.gate.top_k, -1).contiguous().view(-1, 1, self.config.hidden_size) + for i in range(self.gate.top_k): + expert_idx = topk_indices[:, i] + curr_weight = topk_weights[:, i] + gate_qweight = self.all_gate_qweight[expert_idx].reshape( + seq_len * self.out_features_gate, + self.in_features_gate // self.group_size, + (self.group_size * self.bits) // 8, + ) + gate_scales = self.all_gate_scales[expert_idx].reshape( + seq_len * self.out_features_gate * (self.in_features_gate // self.group_size) + ) + gate_qzeros = self.all_gate_qzeros[expert_idx].reshape( + seq_len * self.out_features_gate, self.in_features_gate // self.group_size + ) + gate_gidx = self.all_gate_gidx[expert_idx].reshape(seq_len * self.in_features_gate) + + up_qweight = self.all_up_qweight[expert_idx].reshape( + seq_len * self.out_features_up, + self.in_features_up // self.group_size, + (self.group_size * self.bits) // 8, + ) + up_scales = self.all_up_scales[expert_idx].reshape( + seq_len * self.out_features_up * (self.in_features_up // self.group_size) + ) + up_qzeros = self.all_up_qzeros[expert_idx].reshape( + seq_len * self.out_features_up, self.in_features_up // self.group_size + ) + up_gidx = self.all_up_gidx[expert_idx].reshape(seq_len * self.in_features_up) + + down_qweight = self.all_down_qweight[expert_idx].reshape( + seq_len * self.out_features_down, + self.in_features_down // self.group_size, + (self.group_size * self.bits) // 8, + ) + down_scales = self.all_down_scales[expert_idx].reshape( + seq_len * self.out_features_down * (self.in_features_down // self.group_size) + ) + down_qzeros = self.all_down_qzeros[expert_idx].reshape( + seq_len * self.out_features_down, self.in_features_down // self.group_size + ) + down_gidx = self.all_down_gidx[expert_idx].reshape(seq_len * self.in_features_down) + + gate_out = QuantLinearTorchFunction.apply( + hidden_states, + gate_qweight, + gate_scales, + gate_qzeros, + gate_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_gate, + self.out_features_gate * seq_len, + ) + + up_out = QuantLinearTorchFunction.apply( + hidden_states, + up_qweight, + up_scales, + up_qzeros, + up_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_up, + self.out_features_up * seq_len, + ) + + hidden = self.act_fn(gate_out) * up_out + down_out = QuantLinearTorchFunction.apply( + hidden, + down_qweight, + down_scales, + down_qzeros, + down_gidx if self.act_order else None, + self.bits, + self.group_size, + self.in_features_down, + self.out_features_down, + ) + down_out = down_out.reshape(seq_len, self.out_features_down) + final_hidden_states += down_out * curr_weight.unsqueeze(1) + + return final_hidden_states + + def moe_weights_as_activations(self, hidden_states, router_probs, router_weights): + return QMOE.apply( + hidden_states, + router_weights, + self.fc1_experts_weights, + self.fc1_scales, + self.fc2_experts_weights, + self.fc2_scales, + self.fc3_experts_weights, + self.fc3_scales, + router_probs, + self.config.hidden_act, + self.group_size, + self.bits, + self.num_experts_per_tok, ) - gate_out = torch.bmm(expert_in, gate_proj) - up_out = torch.bmm(expert_in, up_proj) - hidden = self.act_fn(gate_out) * up_out - expert_output = torch.bmm(hidden, down_proj) - experts_out = expert_output.view(seq_len, self.gate.top_k, self.config.hidden_size) - experts_out = experts_out * topk_weights.unsqueeze(-1) - final_hidden_states = torch.einsum("abc->ac", experts_out) + @torch.no_grad() + def original_moe(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + def moe_waa_unpack(self, hidden_states, topk_indices, topk_weights): + gate_proj_unpacked = cast_to_uint4(self.all_gate_qweight) + gate_zeros_unpacked = cast_to_uint4(self.all_gate_qzeros) + gate_proj_dq = dequantize_linear(gate_proj_unpacked, self.all_gate_scales, gate_zeros_unpacked, self.group_size) + + up_proj_unpacked = cast_to_uint4(self.all_up_qweight) + up_zeros_unpacked = cast_to_uint4(self.all_up_qzeros) + up_proj_dq = dequantize_linear(up_proj_unpacked, self.all_up_scales, up_zeros_unpacked, self.group_size) - return final_hidden_states.type(hidden_states.dtype) + down_proj_unpacked = cast_to_uint4(self.all_down_qweight) + down_zeros_unpacked = cast_to_uint4(self.all_down_qzeros) + down_proj_dq = dequantize_linear(down_proj_unpacked, self.all_down_scales, down_zeros_unpacked, self.group_size) + + num_experts = self.all_gate_qweight.shape[0] + expert_in = hidden_states.unsqueeze(0).expand(num_experts, -1, -1) + gate_out = torch.bmm(expert_in, gate_proj_dq.transpose(1, 2).to(expert_in.dtype)) + up_out = torch.bmm(expert_in, up_proj_dq.transpose(1, 2).to(expert_in.dtype)) + hidden = self.act_fn(gate_out) * up_out + down_out = torch.bmm(hidden, down_proj_dq.transpose(1, 2).to(expert_in.dtype)) + + routed_out = down_out.transpose(0, 1) + selected_out = torch.gather( + routed_out, + 1, + topk_indices.unsqueeze(-1).expand(-1, self.gate.top_k, self.out_features_down), + ) + return torch.einsum("abc,ab->ac", selected_out, topk_weights) def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) + topk_indices, topk_weights, router_probs, router_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.moe_waa_unpack(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states @@ -829,83 +1136,333 @@ class QEffPrefillOnlyDeepseekV3MoE(nn.Module): def __qeff_init__( self, ): - for exp in self.experts: - gate_proj = torch.nn.Linear(self.config.hidden_size, self.config.moe_intermediate_size, bias=False) - up_proj = torch.nn.Linear(self.config.hidden_size, self.config.moe_intermediate_size, bias=False) - down_proj = torch.nn.Linear(self.config.moe_intermediate_size, self.config.hidden_size, bias=False) + if hasattr(self, "all_gate_qweight") and not hasattr(self, "experts"): + return + + # Get common parameters from first expert + first_expert = self.experts[0] + self._qeff_quantized_experts = hasattr(first_expert.gate_proj, "bits") + if not self._qeff_quantized_experts: + return + + self.bits = first_expert.gate_proj.bits + self.group_size = first_expert.gate_proj.group_size + self.act_fn = first_expert.act_fn + assert first_expert.gate_proj.act_order == first_expert.up_proj.act_order == first_expert.down_proj.act_order, ( + "act_order mismatch" + ) + self.act_order = first_expert.gate_proj.act_order - gate_proj.weight = torch.nn.Parameter(exp.gate_proj.compressor.decompress_module(exp.gate_proj)) - up_proj.weight = torch.nn.Parameter(exp.up_proj.compressor.decompress_module(exp.up_proj)) - down_proj.weight = torch.nn.Parameter(exp.down_proj.compressor.decompress_module(exp.down_proj)) + # Store dimensions for dequantization + self.in_features_gate, self.out_features_gate = ( + first_expert.gate_proj.in_features, + first_expert.gate_proj.out_features, + ) + self.in_features_up, self.out_features_up = first_expert.up_proj.in_features, first_expert.up_proj.out_features + self.in_features_down, self.out_features_down = ( + first_expert.down_proj.in_features, + first_expert.down_proj.out_features, + ) - setattr(exp, "gate_proj", gate_proj) - setattr(exp, "up_proj", up_proj) - setattr(exp, "down_proj", down_proj) + # Stack all parameters along a new dimension (expert dimension) + self.all_gate_qweight = torch.nn.Parameter( + torch.stack([exp.gate_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // 2 + ), + requires_grad=False, + ) + self.all_gate_scales = torch.nn.Parameter( + torch.stack([exp.gate_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_gate, self.in_features_gate // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + # TODO: Since we know qzeros is always 8 -> Just embed this once into the operator as parameter -> explore this later + self.all_gate_qzeros = torch.nn.Parameter( + torch.stack([exp.gate_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_gate_gidx = torch.nn.Parameter( + torch.stack([exp.gate_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - def moe(self, hidden_states: torch.Tensor, topk_weights: torch.Tensor, expert_mask: torch.Tensor, num_experts: int): - final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - for expert_idx in range(num_experts): - expert = self.experts[expert_idx] - gate_out = expert.gate_proj(hidden_states) - up_out = expert.up_proj(hidden_states) - hidden = expert.act_fn(gate_out) * up_out - expert_output = expert.down_proj(hidden) - current_hidden_states = expert_output * expert_mask[:, expert_idx].unsqueeze(-1) - final_hidden_states += current_hidden_states + self.all_up_qweight = torch.nn.Parameter( + torch.stack([exp.up_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // 2 + ), + requires_grad=False, + ) + self.all_up_scales = torch.nn.Parameter( + torch.stack([exp.up_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_up, self.in_features_up // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + self.all_up_qzeros = torch.nn.Parameter( + torch.stack([exp.up_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_up, self.in_features_up // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_up_gidx = torch.nn.Parameter( + torch.stack([exp.up_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) - return final_hidden_states.type(hidden_states.dtype) + self.all_down_qweight = torch.nn.Parameter( + torch.stack([exp.down_proj.qweight for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // 2 + ), + requires_grad=False, + ) + self.all_down_scales = torch.nn.Parameter( + torch.stack([exp.down_proj.scales for exp in self.experts], dim=0) + .reshape(-1, self.out_features_down, self.in_features_down // self.group_size) + .to(torch.float16), + requires_grad=False, + ) + self.all_down_qzeros = torch.nn.Parameter( + torch.stack([exp.down_proj.qzeros for exp in self.experts], dim=0).reshape( + -1, self.out_features_down, self.in_features_down // (self.group_size * 2) + ), + requires_grad=False, + ) + self.all_down_gidx = torch.nn.Parameter( + torch.stack([exp.down_proj.g_idx for exp in self.experts], dim=0), requires_grad=False + ) + del self.experts - def orig_moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): - r""" - CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused - to not have to do a loop here (deepseek has 256 experts soooo yeah). + def _cumsum_scatter_gather_update_expert_blocked( + self, + x: torch.Tensor, + T2Ei: torch.Tensor, + slot_gate_qweight: torch.Tensor, + slot_gate_scales: torch.Tensor, + slot_gate_qzeros: torch.Tensor, + slot_up_qweight: torch.Tensor, + slot_up_scales: torch.Tensor, + slot_up_qzeros: torch.Tensor, + slot_down_qweight: torch.Tensor, + slot_down_scales: torch.Tensor, + slot_down_qzeros: torch.Tensor, + routing_weight: torch.Tensor, + expert_out: torch.Tensor, + packed_chunk_size: int, + num_q_ffn_blocks: Optional[int] = None, + ) -> torch.Tensor: + """Cumsum-scatter-gather-update expert helper for NSP-blocked dispatch. + + Accumulates one local expert's contribution in-place onto ``expert_out``. + Uses a packed/cumsum layout so the MLP runs only over active rows, then + scatters the weighted output back to original token positions. + + Shapes: + x : [T, H] + T2Ei : [num_nsp, T] (bool) + W_g, W_u : [num_nsp, H, I] + W_d : [num_nsp, I, H] + routing_weight : [num_nsp, T] + expert_out : [num_nsp, T, H] (accumulator, in-out) """ - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) - expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] - mask = expert_mask[expert_idx] - token_indices, weight_indices = torch.where(mask) + batch_size, seq_len = T2Ei.shape + packed_chunk_size = max(1, min(packed_chunk_size, seq_len)) - if token_indices.numel() > 0: - expert_weights = topk_weights[token_indices, weight_indices] - expert_input = hidden_states[token_indices] - expert_output = expert(expert_input) - weighted_output = expert_output * expert_weights.unsqueeze(-1) - final_hidden_states.index_add_(0, token_indices, weighted_output) + if num_q_ffn_blocks is not None: + assert seq_len % num_q_ffn_blocks == 0, "Something went wrong" + packed_chunk_size = seq_len // num_q_ffn_blocks + else: + num_q_ffn_blocks = seq_len // packed_chunk_size + + matched_idx = _build_matched_idx_from_cumsum(T2Ei) + valid_rows = torch.einsum("bi->b", T2Ei.to(torch.int32)).unsqueeze(1) + row_range = torch.arange(packed_chunk_size, dtype=torch.int32, device=x.device).unsqueeze(0) + x_expanded = x.unsqueeze(0).expand(batch_size, -1, -1) + + for chunk_idx in range(num_q_ffn_blocks): + print("executing chunk", chunk_idx) + packed_start = chunk_idx * packed_chunk_size + if chunk_idx == num_q_ffn_blocks - 1: + packed_stop = seq_len + else: + packed_stop = packed_start + packed_chunk_size - # in original deepseek, the output of the experts are gathered once we leave this module - # thus the moe module is itelsf an IsolatedParallel module - # and all expert are "local" meaning we shard but we don't gather - return final_hidden_states.type(hidden_states.dtype) + chunk_matched_idx = matched_idx[:, packed_start:packed_stop] - def forward(self, hidden_states): - """ - Forward pass of MoE block. - """ - residuals = hidden_states - orig_shape = hidden_states.shape - topk_indices, topk_weights = self.gate(hidden_states) - # orig_out = self.orig_moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + x_chunk = CtxGatherFunc3DGeneralized.apply(x_expanded, chunk_matched_idx) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - mask = torch.zeros(hidden_states.shape[0], self.config.n_routed_experts) - mask.scatter_(1, topk_indices, topk_weights) - if os.environ.get("NUM_FFN_BLOCKS", None) is not None and os.environ.get("FFN_W_BLOCK_SIZE", None) is not None: - hidden_states = self.moe_blocked_weights_forward( - hidden_states, topk_weights, mask, self.config.n_routed_experts - ).view(*orig_shape) - elif os.environ.get("NUM_FFN_BLOCKS", None) is not None: - hidden_states = self.moe_blocked_forward( - hidden_states, topk_weights, mask, self.config.n_routed_experts - ).view(*orig_shape) - else: - hidden_states = self.moe(hidden_states, topk_weights, mask, self.config.n_routed_experts).view(*orig_shape) + gate_proj_unpacked = cast_to_uint4(slot_gate_qweight) + gate_zeros_unpacked = cast_to_uint4(slot_gate_qzeros) + gate_proj_dq = dequantize_linear(gate_proj_unpacked, slot_gate_scales, gate_zeros_unpacked, self.group_size) - hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states + up_proj_unpacked = cast_to_uint4(slot_up_qweight) + up_zeros_unpacked = cast_to_uint4(slot_up_qzeros) + up_proj_dq = dequantize_linear(up_proj_unpacked, slot_up_scales, up_zeros_unpacked, self.group_size) + + down_proj_unpacked = cast_to_uint4(slot_down_qweight) + down_zeros_unpacked = cast_to_uint4(slot_down_qzeros) + + down_proj_dq = dequantize_linear(down_proj_unpacked, slot_down_scales, down_zeros_unpacked, self.group_size) + + gate_out = torch.bmm(x_chunk, gate_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + up_out = torch.bmm(x_chunk, up_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + hidden = self.act_fn(gate_out) * up_out + down_out = torch.bmm(hidden, down_proj_dq.transpose(1, 2).to(x_chunk.dtype)) + + rw_chunk = CtxGatherFunc3DGeneralized.apply(routing_weight, chunk_matched_idx) + old_expert_out = CtxGatherFunc3DGeneralized.apply(expert_out, chunk_matched_idx) + chunk_valid_rows = torch.clamp(valid_rows - packed_start, min=0, max=packed_chunk_size) + current_expert_out = ( + torch.where( + (row_range < chunk_valid_rows).unsqueeze(-1), + down_out, + torch.zeros_like(down_out), + ) + * rw_chunk + ) + updated_chunk = old_expert_out + current_expert_out + expert_out = CtxScatterFunc3DGeneralized.apply(expert_out, chunk_matched_idx, updated_chunk) + + return expert_out + + def _forward_expert_blocked( + self, + x: torch.Tensor, + local_T2E: torch.Tensor, + routing_weights: torch.Tensor, + num_q_ffn_blocks: Optional[int] = None, + ) -> torch.Tensor: + T, H = x.shape + num_nsp = EXPERT_BLOCKING_NUM_NSP + if len(self.experts) % num_nsp != 0: + raise ValueError( + f"num_experts ({len(self.experts)}) must be divisible by EXPERT_BLOCKING_NUM_NSP ({num_nsp})" + ) + local_experts = len(self.experts) // num_nsp + routing_weights_unsqueezed = routing_weights.unsqueeze(-1) + expert_out = x.new_zeros((num_nsp, T, H)) + + local_gate_qweight = ( + self.all_gate_qweight.view(local_experts, num_nsp, self.out_features_gate, self.in_features_gate // 2) + .transpose(0, 1) + .contiguous() + ) + local_gate_scales = ( + self.all_gate_scales.view( + local_experts, num_nsp, self.out_features_gate, self.in_features_gate // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_gate_qzeros = ( + self.all_gate_qzeros.view( + local_experts, num_nsp, self.out_features_gate, self.in_features_gate // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + + local_up_qweight = ( + self.all_up_qweight.view(local_experts, num_nsp, self.out_features_up, self.in_features_up // 2) + .transpose(0, 1) + .contiguous() + ) + local_up_scales = ( + self.all_up_scales.view( + local_experts, num_nsp, self.out_features_up, self.in_features_up // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_up_qzeros = ( + self.all_up_qzeros.view( + local_experts, num_nsp, self.out_features_up, self.in_features_up // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + + local_down_qweight = ( + self.all_down_qweight.view(local_experts, num_nsp, self.out_features_down, self.in_features_down // 2) + .transpose(0, 1) + .contiguous() + ) + local_down_scales = ( + self.all_down_scales.view( + local_experts, num_nsp, self.out_features_down, self.in_features_down // self.group_size + ) + .transpose(0, 1) + .contiguous() + ) + local_down_qzeros = ( + self.all_down_qzeros.view( + local_experts, num_nsp, self.out_features_down, self.in_features_down // (self.group_size * 2) + ) + .transpose(0, 1) + .contiguous() + ) + for slot in range(local_experts): + print(f"executing slot {slot}") + T2Ei = local_T2E[:, slot, :] + expert_out = self._cumsum_scatter_gather_update_expert_blocked( + x=x, + T2Ei=T2Ei, + slot_gate_qweight=local_gate_qweight[:, slot], + slot_gate_scales=local_gate_scales[:, slot], + slot_gate_qzeros=local_gate_qzeros[:, slot], + slot_up_qweight=local_up_qweight[:, slot], + slot_up_scales=local_up_scales[:, slot], + slot_up_qzeros=local_up_qzeros[:, slot], + slot_down_qweight=local_down_qweight[:, slot], + slot_down_scales=local_down_scales[:, slot], + slot_down_qzeros=local_down_qzeros[:, slot], + routing_weight=routing_weights_unsqueezed[:, slot], + expert_out=expert_out, + packed_chunk_size=EXPERT_BLOCKING_PACKED_CHUNK_SIZE, + num_q_ffn_blocks=num_q_ffn_blocks, + ) + return torch.einsum("bij->ij", expert_out) + + def forward( + self, hidden_states: torch.Tensor, num_q_ffn_blocks: Optional[int] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + topk_idx, topk_weight, _, _ = self.gate(hidden_states) + B, S, H = hidden_states.shape + T = B * S + x = hidden_states.view(T, H) + + if len(self.experts) % EXPERT_BLOCKING_NUM_NSP == 0: + expert_ids = torch.arange( + len(self.experts) // EXPERT_BLOCKING_NUM_NSP, + device=x.device, + dtype=topk_idx.dtype, + ).unsqueeze(0) * EXPERT_BLOCKING_NUM_NSP + torch.arange( + EXPERT_BLOCKING_NUM_NSP, device=x.device, dtype=topk_idx.dtype + ).unsqueeze(1) # [N, L] + eq = topk_idx.unsqueeze(0).unsqueeze(0) == expert_ids.unsqueeze(-1).unsqueeze(-1) + local_T2E = eq.to(topk_idx.dtype).sum(dim=-1) > 0 # [N, L, T] + routing_weights = (eq.to(topk_weight.dtype) * topk_weight.unsqueeze(0).unsqueeze(0)).sum(dim=-1) + + expert_out = self._forward_expert_blocked( + x=x, local_T2E=local_T2E, routing_weights=routing_weights, num_q_ffn_blocks=num_q_ffn_blocks + ) + self.shared_experts(hidden_states) + return expert_out.view(B, S, H) + + routing_weights = torch.zeros(T, self.config.n_routed_experts, device=x.device, dtype=topk_weight.dtype) + routing_weights.scatter_(1, topk_idx, topk_weight) + + final_hidden_states = x.new_zeros((T, H)) + for expert_idx in range(self.config.n_routed_experts): + expert = self.experts[expert_idx] + gate_out = expert.gate_proj(hidden_states) + up_out = expert.up_proj(hidden_states) + hidden = expert.act_fn(gate_out) * up_out + expert_output = expert.down_proj(hidden) + current_hidden_states = expert_output * routing_weights[:, expert_idx].unsqueeze(-1) + final_hidden_states = final_hidden_states.view(B, S, H) + current_hidden_states + + final_hidden_states = final_hidden_states + self.shared_experts(hidden_states) + return final_hidden_states.view(B, S, H) class QEffDeepseekV3DecoderLayer(nn.Module): @@ -924,6 +1481,9 @@ def forward( cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, mla_absorption: Optional[Dict[str, bool]] = None, + num_q_ffn_blocks: Optional[int] = None, + sin_cached=None, + cos_cached=None, **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states @@ -945,6 +1505,8 @@ def forward( use_cache=use_cache, cache_position=cache_position, mla_absorption=mla_absorption, + sin_cached=sin_cached, + cos_cached=cos_cached, **kwargs, ) else: @@ -958,13 +1520,18 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + sin_cached=sin_cached, + cos_cached=cos_cached, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + if num_q_ffn_blocks is not None and self.mlp.__class__.__name__ == "DeepseekV3MoE": + hidden_states = self.mlp(hidden_states, num_q_ffn_blocks) + else: + hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) @@ -1003,6 +1570,8 @@ def __qeff_init__(self): base=self.config.rope_theta, **kwargs, ) + self.sin_cached = torch.nn.Parameter(self.rotary_emb.sin_cached.to(torch.float16)) + self.cos_cached = torch.nn.Parameter(self.rotary_emb.cos_cached.to(torch.float16)) def forward( self, @@ -1065,6 +1634,9 @@ def forward( all_self_attns = () if output_attentions else None next_decoder_cache = None + sin = self.sin_cached[position_ids].unsqueeze(1) + cos = self.cos_cached[position_ids].unsqueeze(1) + num_q_ffn_blocks = getattr(self, "num_q_blocks_ffn", None) for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) @@ -1081,6 +1653,9 @@ def forward( cache_position=cache_position, position_embeddings=position_embeddings, mla_absorption=mla_absorption, + num_q_ffn_blocks=num_q_ffn_blocks, + sin_cached=sin, + cos_cached=cos, **kwargs, ) diff --git a/QEfficient/transformers/models/kimi_k25/__init__.py b/QEfficient/transformers/models/kimi_k25/__init__.py new file mode 100644 index 0000000000..da26921c50 --- /dev/null +++ b/QEfficient/transformers/models/kimi_k25/__init__.py @@ -0,0 +1,7 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py new file mode 100644 index 0000000000..a083018308 --- /dev/null +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -0,0 +1,1153 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import math +from typing import List, Optional, Tuple, Type, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import activations + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.cache_utils import Cache +from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast + +from QEfficient.utils import constants + + +def eager_attention_forward(q, k, v, **kwargs): + q_cu_seqlens = kwargs.get("q_cu_seqlens") + seq_length = q.shape[0] + if q_cu_seqlens is not None: + attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) + for idx in range(1, len(q_cu_seqlens)): + attention_mask[ + ..., + q_cu_seqlens[idx - 1] : q_cu_seqlens[idx], + q_cu_seqlens[idx - 1] : q_cu_seqlens[idx], + ] = True + else: + attention_mask = None + + q = q.transpose(0, 1) # (num_heads, seq_len, head_dim) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + if attention_mask is not None: + attn_weight = attn_weight + attention_mask + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) + attn_out = attn_weight @ v + attn_out = attn_out.transpose(0, 1) # (seq_len, num_heads, head_dim) + return attn_out.reshape(attn_out.shape[0], -1) + + +VL_VISION_ATTENTION_FUNCTIONS = {"eager": eager_attention_forward} + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate +def get_rope_shape(org, interpolation_mode, shape): + return ( + F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) + + +def apply_rope(xq, xk, freqs_cis): + # Support both complex freqs (..., dim//2) and stacked real/imag (2, ..., dim//2) + if torch.is_complex(freqs_cis): + freqs_cis = freqs_cis.unsqueeze(-2) + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + # freqs_cis shape: (2, seq_len, dim//2) + # xq shape: (seq_len, num_heads, head_dim) + freqs_cos = freqs_cis[0].unsqueeze(-2) # (seq_len, 1, dim//2) + freqs_sin = freqs_cis[1].unsqueeze(-2) # (seq_len, 1, dim//2) + xq_r = xq.float().view(*xq.shape[:-1], -1, 2) + xq_r0, xq_r1 = xq_r[..., 0], xq_r[..., 1] + xq_out_r = xq_r0 * freqs_cos - xq_r1 * freqs_sin + xq_out_i = xq_r0 * freqs_sin + xq_r1 * freqs_cos + xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2) + xk_r = xk.float().view(*xk.shape[:-1], -1, 2) + xk_r0, xk_r1 = xk_r[..., 0], xk_r[..., 1] + xk_out_r = xk_r0 * freqs_cos - xk_r1 * freqs_sin + xk_out_i = xk_r0 * freqs_sin + xk_r1 * freqs_cos + xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +class QEffLearnable2DInterpPosEmbDivided_fixed(nn.Module): + def __qeff_init__(self): + self.interpolation_mode = "bilinear" + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, "dim must be divisible by 4" + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" + + def _ensure_precomputed_freqs(self, device: torch.device) -> None: + if not hasattr(self, "freqs_cis"): + self.register_buffer("freqs_cis", self._precompute_freqs_cis(device), persistent=False) + elif self.freqs_cis.device != device: + self.freqs_cis = self._precompute_freqs_cis(device) + + if not hasattr(self, "freqs_cos"): + self.register_buffer("freqs_cos", self.freqs_cis.real.contiguous(), persistent=False) + elif self.freqs_cos.device != device: + self.freqs_cos = self.freqs_cis.real.contiguous() + + if not hasattr(self, "freqs_sin"): + self.register_buffer("freqs_sin", self.freqs_cis.imag.contiguous(), persistent=False) + elif self.freqs_sin.device != device: + self.freqs_sin = self.freqs_cis.imag.contiguous() + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) # C/4 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis(self, grid_thws: torch.Tensor, device: torch.device) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + self._ensure_precomputed_freqs(device) + + shapes = grid_thws.tolist() + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) for t, h, w in shapes], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + *, + attn_implementation: str = "flash_attention_2", + activation=F.gelu, + attn_bias: bool = False, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation) + self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func( + xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn, + ) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class QEffMoonViT3dEncoder(nn.Module): + def __qeff_init__(self): + if not hasattr(self, "blocks") or len(self.blocks) == 0: + return + + old_blocks = list(self.blocks) + first_block = old_blocks[0] + self.block_cfg = { + "num_heads": first_block.num_heads, + "hidden_dim": first_block.hidden_dim, + "mlp_dim": first_block.mlp.fc0.out_features, + "activation": PytorchGELUTanh(), + "attn_bias": first_block.wqkv.bias is not None, + "attn_implementation": "eager", + } + + head_dim = first_block.hidden_size_per_attention_head + max_height = getattr(self.rope_2d, "max_height", 512) + max_width = getattr(self.rope_2d, "max_width", 512) + theta_base = getattr(self.rope_2d, "theta_base", 10000) + self.rope_2d = Rope2DPosEmbRepeated(head_dim, max_height, max_width, theta_base=theta_base) + + new_blocks = [] + for old_block in old_blocks: + old_weight = old_block.wqkv.weight + if old_weight.is_meta: + with torch.device("meta"): + new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + new_block.load_state_dict(old_block.state_dict(), assign=True) + else: + new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + new_block.load_state_dict(old_block.state_dict()) + new_blocks.append(new_block.to(device=old_weight.device, dtype=old_weight.dtype)) + self.blocks = nn.ModuleList(new_blocks) + + +class QEffKimiK25EncoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.config = self.model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.model.vision_tower.encoder.blocks[0].__class__} + + def forward(self, pixel_values: torch.Tensor, h_shape: torch.Tensor, w_shape: torch.Tensor) -> torch.Tensor: + """ + ONNX-exportable forward that runs only the vision tower and mm_projector. + Uses h_shape and w_shape (int64 ones tensors of length h and w respectively) + to derive spatial dimensions via .shape[0], enabling dynamic axis export. + + Args: + pixel_values: Preprocessed image pixel values (num_patches, channels*patch_h*patch_w). + h_shape: int64 ones tensor of shape (h,) encoding number of patch rows. + w_shape: int64 ones tensor of shape (w,) encoding number of patch columns. + + Returns: + vision_embeds: Projected image embeddings as a single tensor. For multiple + same-sized images, embeddings are concatenated in input order. + """ + + h_shape = h_shape.to(pixel_values.device) + w_shape = w_shape.to(pixel_values.device) + + # Keep them in ONNX graph + dummy = (h_shape.float().sum() + w_shape.float().sum()) * 0.0 + pixel_values = pixel_values + dummy + + h = h_shape.shape[0] + w = w_shape.shape[0] + num_tokens_per_image = h * w + if pixel_values.shape[0] % num_tokens_per_image != 0: + raise ValueError( + f"pixel_values first dimension ({pixel_values.shape[0]}) must be divisible by h*w ({num_tokens_per_image})." + ) + num_images = pixel_values.shape[0] // num_tokens_per_image + + target_dtype = self.model.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + hidden_states = self.model.vision_tower.patch_embed.proj(pixel_values).view(pixel_values.size(0), -1) + + pos_emb_module = self.model.vision_tower.patch_embed.pos_emb + pos_emb_2d = get_rope_shape( + pos_emb_module.weight, + interpolation_mode=pos_emb_module.interpolation_mode, + shape=(h, w), + ) + hidden_states = hidden_states + pos_emb_2d.repeat(num_images, 1) + + rope_2d = self.model.vision_tower.encoder.rope_2d + rope_2d._ensure_precomputed_freqs(pixel_values.device) + freqs_cos = rope_2d.freqs_cos[:h, :w].reshape(-1, rope_2d.dim // 2).repeat(num_images, 1) + freqs_sin = rope_2d.freqs_sin[:h, :w].reshape(-1, rope_2d.dim // 2).repeat(num_images, 1) + freqs_cis = torch.stack([freqs_cos, freqs_sin], dim=0) + + if num_images == 1: + cu_seqlens = None + else: + lengths = torch.full( + (num_images + 1,), num_tokens_per_image, dtype=torch.int64, device=hidden_states.device + ) + lengths[0] = 0 + cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int64) + max_seqlen = num_tokens_per_image + + encoder_dtype = self.model.vision_tower.encoder.blocks[0].wqkv.weight.dtype + hidden_states = hidden_states.to(encoder_dtype) + freqs_cis = freqs_cis.to(encoder_dtype) + for block in self.model.vision_tower.encoder.blocks: + hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=freqs_cis) + + final_ln_dtype = self.model.vision_tower.encoder.final_layernorm.weight.dtype + hidden_states = self.model.vision_tower.encoder.final_layernorm(hidden_states.to(final_ln_dtype)) + + merge_kernel_size = self.model.vision_tower.merge_kernel_size + kernel_height, kernel_width = merge_kernel_size + new_height = h // kernel_height + new_width = w // kernel_width + d_model = hidden_states.size(-1) + reshaped = hidden_states.view(num_images, new_height, kernel_height, new_width, kernel_width, d_model) + reshaped = reshaped.permute(0, 1, 3, 2, 4, 5).contiguous() + merged = reshaped.view(num_images * new_height * new_width, kernel_height * kernel_width, -1) + + pre_norm_dtype = self.model.mm_projector.pre_norm.weight.dtype + merged = merged.to(pre_norm_dtype) + vision_embeds = self.model.mm_projector.proj(self.model.mm_projector.pre_norm(merged).view(merged.shape[0], -1)) + return vision_embeds + + +class QEffKimiK25DecoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.language_model = self.model.language_model + self.lm_head = self.language_model.lm_head + self.config = self.model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + """ + Return the set of class used as the repeated layer across the model for subfunction extraction. + Notes: + This method should return the *class object* (not an instance). + Downstream code can use this to find/build subfunctions for repeated blocks. + """ + return {self.language_model.model.layers[0].__class__} + + def forward( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + vision_embeds: Optional[torch.FloatTensor] = None, + image_idx: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + batch_index: Optional[torch.LongTensor] = None, + comp_ctx_lengths: Optional[List[int]] = None, + use_cache: Optional[bool] = None, + **kwargs, + ) -> Tuple: + if inputs_embeds is None: + inputs_embeds = self.model.get_input_embeddings()(input_ids) + else: + inputs_embeds = inputs_embeds + + vision_embeds_for_state = None + if vision_embeds is not None: + if vision_embeds.dim() == 3: + if vision_embeds.shape[0] != 1: + raise ValueError( + f"Expected vision_embeds batch dim to be 1, got shape {tuple(vision_embeds.shape)}" + ) + vision_embeds_for_state = vision_embeds[0].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + elif vision_embeds.dim() == 2: + vision_embeds_for_state = vision_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + else: + raise ValueError(f"Expected vision_embeds rank 2 or 3, got {vision_embeds.dim()}.") + + if vision_embeds_for_state is not None and input_ids is not None: + if attention_mask is None: + if position_ids is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) + else: + attention_mask = (position_ids >= 0).to(dtype=torch.long, device=input_ids.device) + if image_idx is None: + image_idx = torch.zeros((input_ids.shape[0], 1), dtype=torch.int64, device=input_ids.device) + selected = input_ids == self.config.media_placeholder_token_id + selected_any = selected.any(dim=1, keepdim=True) + selected_image_tokens = selected.to(torch.int64).sum(dim=1, keepdim=True) + inputs_embeds = inputs_embeds.to(vision_embeds_for_state.dtype) + + merged_inputs_embeds, merged_attention_mask, _, merged_position_ids = ( + self.model._qeff_merge_input_ids_with_image_features( + image_features=vision_embeds_for_state, + inputs_embeds=inputs_embeds, + input_ids=input_ids, + attention_mask=attention_mask, + labels=None, + ) + ) + + inputs_embeds = merged_inputs_embeds + attention_mask = merged_attention_mask + merged_image_tokens = torch.full( + (1, 1), vision_embeds_for_state.shape[0], device=image_idx.device, dtype=torch.int64 + ) + default_image_idx = torch.clamp(merged_image_tokens - 1, min=0) + input_batch = torch.full((1, 1), input_ids.shape[0], device=image_idx.device, dtype=torch.int64) + image_idx_batch = torch.full((1, 1), image_idx.shape[0], device=image_idx.device, dtype=torch.int64) + use_default_image_idx = torch.logical_and(torch.logical_not(selected_any), input_batch != image_idx_batch) + effective_image_idx = torch.where(use_default_image_idx, default_image_idx, image_idx) + if position_ids is None: + position_ids = merged_position_ids + else: + # Preserve caller-provided absolute position offset (needed for + # chunked prefill/decode parity) while using merged sequence + # positions for image-expanded tokens. + position_offset = position_ids[:, :1] + effective_image_idx.to( + device=position_ids.device, dtype=position_ids.dtype + ) + position_ids = torch.where( + merged_attention_mask > 0, + merged_position_ids + position_offset, + torch.full_like(merged_position_ids, -1), + ) + + image_position_delta = torch.clamp(merged_image_tokens - selected_image_tokens, min=0) + image_idx = effective_image_idx + selected_any.to(torch.int64) * image_position_delta + + if position_ids is None and attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) + elif position_ids is not None and attention_mask is not None: + position_ids = torch.where(attention_mask > 0, position_ids, torch.full_like(position_ids, -1)) + + outputs = self.model.language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + compressed_kvs=compressed_kvs, + past_key_values=past_key_values, + batch_index=batch_index, + comp_ctx_lengths=comp_ctx_lengths, + use_cache=True if use_cache is None else use_cache, + **kwargs, + ) + logits = outputs[0].float() + + mla_absorption = getattr(self.model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + output_kvs = getattr(outputs, "compressed_kvs", None) + else: + output_kvs = getattr(outputs, "past_key_values", None) + if vision_embeds_for_state is not None: + vision_embeds_for_state = vision_embeds_for_state + torch.zeros_like(vision_embeds_for_state) + image_idx_output = image_idx[:1, :1] if image_idx is not None else image_idx + return logits, vision_embeds_for_state, image_idx_output, output_kvs + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class QEffKimiK25ForConditionalGeneration(nn.Module): + def get_qeff_vision_encoder(self): + return QEffKimiK25EncoderWrapper(self) + + def get_qeff_language_decoder(self): + return QEffKimiK25DecoderWrapper(self) + + def _qeff_merge_input_ids_with_image_features( + self, + image_features: torch.Tensor, + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + + target_device = inputs_embeds.device + image_features = image_features.to(target_device) + num_image_tokens = torch.full((1, 1), image_features.shape[0], device=input_ids.device, dtype=input_ids.dtype) + + image_token_mask = input_ids == image_token_index + non_image_mask = ~image_token_mask + image_token_count = image_token_mask.to(input_ids.dtype).sum(dim=1, keepdim=True) + has_image = image_token_count > 0 + + token_occupation = torch.where( + image_token_mask, + num_image_tokens.view(1, 1).expand_as(input_ids), + torch.ones_like(input_ids), + ) + new_token_positions = torch.cumsum(token_occupation, dim=1) - 1 + + final_seq_len = input_ids.shape[1] + image_features.shape[0] - 1 + merged_positions = torch.arange(final_seq_len, device=input_ids.device, dtype=input_ids.dtype).view(1, 1, -1) + text_position_one_hot = merged_positions == new_token_positions.unsqueeze(-1) + text_position_one_hot = torch.logical_and(text_position_one_hot, non_image_mask.unsqueeze(-1)) + + final_embedding = ( + text_position_one_hot.to(inputs_embeds.dtype).unsqueeze(-1) * inputs_embeds.unsqueeze(2) + ).sum(dim=1) + final_attention_mask = (text_position_one_hot.to(attention_mask.dtype) * attention_mask.unsqueeze(-1)).sum( + dim=1 + ) + + image_features_for_batch = image_features.unsqueeze(0).expand(input_ids.shape[0], -1, -1) + media_embedding = torch.cat([image_features_for_batch, final_embedding[:, image_features.shape[0] :, :]], dim=1) + media_attention_mask = torch.cat( + [ + torch.ones( + (input_ids.shape[0], image_features.shape[0]), + dtype=attention_mask.dtype, + device=input_ids.device, + ), + final_attention_mask[:, image_features.shape[0] :], + ], + dim=1, + ) + final_embedding = torch.where(has_image.unsqueeze(-1), media_embedding, final_embedding) + final_attention_mask = torch.where(has_image, media_attention_mask, final_attention_mask) + + position_ids = torch.cumsum(final_attention_mask, dim=1) - 1 + position_ids = torch.where(final_attention_mask == 0, torch.full_like(position_ids, -1), position_ids) + + pad_token_mask = input_ids == pad_token_id + pad_position_one_hot = torch.logical_and( + merged_positions == new_token_positions.unsqueeze(-1), pad_token_mask.unsqueeze(-1) + ) + pad_positions_mask = pad_position_one_hot.any(dim=1) + final_embedding = torch.where( + pad_positions_mask.unsqueeze(-1), torch.zeros_like(final_embedding), final_embedding + ) + + return final_embedding, final_attention_mask, labels, position_ids + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, + grid_thws: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + compressed_kvs: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and len(pixel_values) > 0 and input_ids.shape[1] != 1: + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to(image_features[0].dtype) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + ) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + compressed_kvs=compressed_kvs, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def get_output_names(self, kv_offload: bool = False): + vision_output_names = ["vision_embeds"] + lang_output_names = ["logits"] + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_output_names.append(f"compressed_kv.{i}_RetainedState") + lang_output_names.append(f"k_pe.{i}_RetainedState") + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_output_names.append(f"past_{kv}.{i}_RetainedState") + + output_names = {} + if kv_offload: + lang_output_names.insert(1, "vision_embeds_RetainedState") + lang_output_names.insert(2, "image_idx_output") + output_names["vision"] = vision_output_names + output_names["lang"] = lang_output_names + else: + lang_output_names.insert(1, "pixel_values_RetainedState") + lang_output_names.insert(2, "image_idx_output") + return lang_output_names + return output_names + + def get_dummy_inputs( + self, + kv_offload: bool = False, + continuous_batching: bool = False, + **kwargs, + ): + prefill_seq_len = kwargs.get("prefill_seq_len") + if prefill_seq_len is None: + prefill_seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + prefill_seq_len = int(prefill_seq_len) + + bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE + fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS + + inputs_shapes = {} + inputs_shapes["pixel_values"] = ( + constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH, + constants.ONNX_EXPORT_IMAGE_DEPTH, + constants.KIMI_PATCH_SIZE, + constants.KIMI_PATCH_SIZE, + ) + inputs_shapes["vision_embeds"] = ( + constants.KIMI_EXAMPLE_IMAGE_NUM_IMAGE_TOKENS, + self.language_model.config.hidden_size, + ) + inputs_shapes["image_idx"] = (1, 1) + inputs_shapes["grid_h"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT + inputs_shapes["grid_w"] = constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH + + vision_inputs = { + "pixel_values": torch.zeros((inputs_shapes["pixel_values"]), dtype=self.config.torch_dtype), + "h_shape": torch.zeros((inputs_shapes["grid_h"]), dtype=torch.int64), + "w_shape": torch.zeros((inputs_shapes["grid_w"]), dtype=torch.int64), + } + + input_ids = torch.zeros((bs, prefill_seq_len), dtype=torch.int64) + input_ids[:, 0] = self.config.media_placeholder_token_id + + lang_inputs = { + "input_ids": input_ids, + "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64).view(1, prefill_seq_len).repeat(bs, 1), + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + pkv_cache = self.language_model.get_dummy_pkv_cache( + config=self.language_model.config, + batch_size=fbs if continuous_batching else bs, + seq_len=constants.ONNX_EXPORT_CTX_LEN, + ) + + if cache_compressed: + lang_inputs["compressed_kvs"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["compressed_kvs"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + else: + lang_inputs["past_key_values"] = [[] for _ in range(self.language_model.config.num_hidden_layers)] + for i in range(self.language_model.config.num_hidden_layers): + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][0].shape, dtype=self.language_model.config.torch_dtype) + ) + lang_inputs["past_key_values"][i].append( + torch.zeros(pkv_cache[0][1].shape, dtype=self.language_model.config.torch_dtype) + ) + + lang_inputs["vision_embeds"] = torch.zeros( + inputs_shapes["vision_embeds"], + dtype=self.language_model.config.torch_dtype, + ) + lang_inputs["image_idx"] = torch.zeros(inputs_shapes["image_idx"], dtype=torch.int64) + + if continuous_batching: + lang_inputs["batch_index"] = torch.arange(bs).view(bs, 1) + + inputs = {} + if kv_offload: + inputs["vision"] = vision_inputs + inputs["lang"] = lang_inputs + else: + lang_inputs.pop("vision_embeds") + inputs = {**vision_inputs, **lang_inputs} + + return inputs + + def get_onnx_dynamic_axes( + self, comp_ctx_lengths: Optional[List[int]] = None, kv_offload: bool = False, continuous_batching: bool = False + ): + vision_dynamic_axes = {} + lang_dynamic_axes = {} + lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} + lang_dynamic_axes["vision_embeds"] = {0: "num_image_tokens"} + if continuous_batching: + lang_dynamic_axes["batch_index"] = {0: "batch_size"} + vision_dynamic_axes = { + "pixel_values": {0: "num_patches"}, + "h_shape": {0: "grid_h"}, + "w_shape": {0: "grid_w"}, + "vision_embeds": {0: "num_image_tokens"}, + } + + mla_absorption = getattr(self.language_model, "mla_absorption", None) + if mla_absorption is not None: + cache_compressed = mla_absorption.get("cache_compressed", False) + else: + cache_compressed = False + + cache_batch_axis = "full_batch_size" if continuous_batching else "batch_size" + + if cache_compressed: + for i in range(self.language_model.config.num_hidden_layers): + lang_dynamic_axes[f"compressed_kv.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} + lang_dynamic_axes[f"k_pe.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} + else: + for i in range(self.language_model.config.num_hidden_layers): + for kv in ["key", "value"]: + lang_dynamic_axes[f"past_{kv}.{i}"] = {0: cache_batch_axis, 2: "ctx_len"} + + if comp_ctx_lengths is not None: + lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} + + dynamic_axes = {} + if kv_offload: + dynamic_axes["vision"] = vision_dynamic_axes + dynamic_axes["lang"] = lang_dynamic_axes + else: + lang_dynamic_axes.pop("vision_embeds") + dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} + return dynamic_axes + + def get_specializations( + self, + batch_size: int, + prefill_seq_len: int, + ctx_len: int, + kv_offload: bool = False, + continuous_batching: bool = False, + kv_cache_batch_size: Optional[int] = None, + full_batch_size: Optional[int] = None, + **compiler_options, + ): + comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) + comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + compiler_options.pop("img_size", None) + image_height = compiler_options.pop("image_height", None) + image_width = compiler_options.pop("image_width", None) + unsupported_shape_args = { + name: compiler_options.pop(name, None) + for name in ("height", "width", "h", "w", "num_patches", "num_image_tokens") + } + if any(value is not None for value in unsupported_shape_args.values()): + raise ValueError( + "Kimi-K2.5 compile expects image_height and image_width. " + "Do not pass height/width, h/w grid dimensions, num_patches, or num_image_tokens." + ) + num_frames = compiler_options.pop("num_frames", 1) + mm_processor_kwargs = compiler_options.pop("mm_processor_kwargs", None) or {} + + prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN + + def normalize_list(value, default): + if value is None: + return [default] + if isinstance(value, int): + return [value] + return list(value) + + def normalize_sized_list(value, count, name): + if value is None: + return None + if isinstance(value, int): + return [value] * count + value = list(value) + if len(value) != count: + raise ValueError(f"Expected {name} to contain {count} entries, got {len(value)}.") + return value + + def validate_dimension_lists(heights, widths, height_name, width_name): + if len(heights) != len(widths): + raise ValueError( + f"Expected {height_name} and {width_name} to contain the same number of entries, " + f"got {len(heights)} and {len(widths)}." + ) + if any(height <= 0 for height in heights) or any(width <= 0 for width in widths): + raise ValueError(f"Expected {height_name} and {width_name} to contain positive values.") + + patch_size = getattr(self.config.vision_config, "patch_size", constants.KIMI_PATCH_SIZE) + merge_kernel_size = getattr(self.config.vision_config, "merge_kernel_size", (2, 2)) + if isinstance(merge_kernel_size, int): + kernel_height = kernel_width = merge_kernel_size + merge_kernel_size = (merge_kernel_size, merge_kernel_size) + else: + kernel_height, kernel_width = merge_kernel_size + + if (image_height is None) != (image_width is None): + raise ValueError("Kimi-K2.5 compile expects image_height and image_width to be provided together.") + + if image_height is not None: + pixel_heights = normalize_list(image_height, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT * patch_size) + pixel_widths = normalize_list(image_width, constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH * patch_size) + validate_dimension_lists(pixel_heights, pixel_widths, "image_height", "image_width") + + in_patch_limit = mm_processor_kwargs.get("in_patch_limit", 16384) + patch_limit_on_one_side = mm_processor_kwargs.get("patch_limit_on_one_side", 512) + factor_height = kernel_height * patch_size + factor_width = kernel_width * patch_size + grid_heights = [] + grid_widths = [] + for pixel_height, pixel_width in zip(pixel_heights, pixel_widths): + scale = min( + 1.0, + math.sqrt( + in_patch_limit / (max(1.0, pixel_width // patch_size) * max(1.0, pixel_height // patch_size)) + ), + patch_limit_on_one_side * patch_size / pixel_width, + patch_limit_on_one_side * patch_size / pixel_height, + ) + resized_height = min(max(1, int(pixel_height * scale)), patch_limit_on_one_side * patch_size) + resized_width = min(max(1, int(pixel_width * scale)), patch_limit_on_one_side * patch_size) + pad_height = (factor_height - resized_height % factor_height) % factor_height + pad_width = (factor_width - resized_width % factor_width) % factor_width + grid_heights.append((resized_height + pad_height) // patch_size) + grid_widths.append((resized_width + pad_width) // patch_size) + else: + grid_heights = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT] + grid_widths = [constants.KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH] + + num_frames = normalize_sized_list(1 if num_frames is None else num_frames, len(grid_heights), "num_frames") + + vision = [] + max_num_image_tokens = 0 + for index, (grid_height, grid_width, frames) in enumerate(zip(grid_heights, grid_widths, num_frames)): + if grid_height % kernel_height != 0 or grid_width % kernel_width != 0: + raise ValueError( + f"Kimi image grid h={grid_height}, w={grid_width} must be divisible by " + f"merge_kernel_size={merge_kernel_size}." + ) + + computed_num_patches = grid_height * grid_width * frames + computed_num_image_tokens = (grid_height // kernel_height) * (grid_width // kernel_width) * frames + max_num_image_tokens = max(max_num_image_tokens, computed_num_image_tokens) + + vision.append( + { + "num_patches": computed_num_patches, + "grid_h": grid_height, + "grid_w": grid_width, + "num_image_tokens": computed_num_image_tokens, + } + ) + + if comp_ctx_lengths_prefill is not None: + lang = [] + + for i in range(0, len(comp_ctx_lengths_prefill)): + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": max_num_image_tokens, + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang.append(lang_prefill) + + for i in range(0, len(comp_ctx_lengths_decode)): + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": "1", + "ctx_len": ctx_len, + "num_image_tokens": max_num_image_tokens, + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang.append(lang_decode) + + else: + lang_prefill = { + "batch_size": 1 if continuous_batching else batch_size, + "seq_len": prefill_seq_len, + "ctx_len": ctx_len, + "num_image_tokens": max_num_image_tokens, + } + if continuous_batching: + lang_prefill["full_batch_size"] = kv_cache_batch_size + else: + lang_prefill["batch_size"] = kv_cache_batch_size + if full_batch_size: + lang_prefill["full_batch_exec_size"] = full_batch_size + + lang_decode = { + "batch_size": full_batch_size if continuous_batching else batch_size, + "seq_len": 1, + "ctx_len": ctx_len, + "num_image_tokens": max_num_image_tokens, + } + + if continuous_batching: + lang_decode["full_batch_size"] = kv_cache_batch_size + else: + lang_decode["batch_size"] = kv_cache_batch_size + + lang = [lang_prefill, lang_decode] + + specializations = {} + + if kv_offload: + specializations["vision"] = vision + specializations["lang"] = lang + return specializations, compiler_options + else: + return [{**vision_spec, **lang_spec} for vision_spec in vision for lang_spec in lang], compiler_options diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 23bc63c73d..dfaa687a0e 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -32,6 +32,12 @@ from QEfficient.base.modeling_qeff import QEFFBaseModel from QEfficient.base.onnx_transforms import FP16ClipTransform, SplitTensorsTransform from QEfficient.base.pytorch_transforms import SplitGateUpWeightsTransform +from QEfficient.exporter.weight_free.transforms import ( + DtypeConversionCheckpointTransform, + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, +) from QEfficient.generation.cloud_infer import QAICInferenceSession, is_retained_state_name from QEfficient.generation.text_generation_inference import ( CloudAI100ExecInfoNew, @@ -72,6 +78,7 @@ FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) from QEfficient.utils import ( apply_kv_cache_prefix, @@ -1068,14 +1075,23 @@ class QEffVisionEncoderForTextImageToTextModel(QEFFBaseModel): of multimodal models for optimal performance on Cloud AI 100 hardware. """ + _hf_auto_class = AutoModelForImageTextToText + _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, CustomOpsTransform, KVCacheTransform, KVCacheExternalModuleMapperTransform, ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def __init__(self, model: nn.modules, **kwargs): """ @@ -1124,6 +1140,8 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=kwargs.get("dynamo", False), + use_weight_free_export=kwargs.get("use_weight_free_export", False), ) def compile( @@ -1202,9 +1220,12 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): of multimodal models for optimal performance on Cloud AI 100 hardware. """ + _hf_auto_class = AutoModelForImageTextToText + _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, FP8BlockWiseDequantQwen3VLMoeTextExpertsToQwen3VLMoeTextExpertsTransform, FP8BlockWiseDequantLinearToLinearTransform, CustomOpsTransform, @@ -1213,6 +1234,12 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): SplitGateUpWeightsTransform, ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): """ @@ -1234,6 +1261,10 @@ def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): 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, @@ -1321,6 +1352,8 @@ def export( export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=kwargs.get("dynamo", False), + use_weight_free_export=kwargs.get("use_weight_free_export", False), ) def compile( @@ -1436,6 +1469,7 @@ def __init__( self.ccl_enabled = False if qaic_config: self.ccl_enabled = 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 # ---Sampling--- @@ -1537,6 +1571,7 @@ def export( layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, offload_pt_weights: Optional[bool] = None, + use_weight_free_export: bool = False, **kwargs, ) -> str: """ @@ -1571,6 +1606,7 @@ def export( enable_chunking=enable_chunking, layerwise_window_size=layerwise_window_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **kwargs, ) bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE @@ -1652,6 +1688,8 @@ def export( export_dir=export_dir, offload_pt_weights=False, use_onnx_subfunctions=use_onnx_subfunctions, + dynamo=use_weight_free_export, + use_weight_free_export=use_weight_free_export, ) # TODO: remove the current pt weight offload capability once CustomLoader is in place @@ -1676,6 +1714,8 @@ def export( prefill_seq_len=prefill_seq_len, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + dynamo=use_weight_free_export, + use_weight_free_export=use_weight_free_export, ) return self.onnx_path @@ -1860,6 +1900,7 @@ def compile( layerwise: bool = False, layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, + use_weight_free_export: bool = False, moe_prefill_packed_chunk_size: int = constants.MOE_PREFILL_PACKED_CHUNK_SIZE, **compiler_options, ) -> str: @@ -1948,6 +1989,7 @@ def compile( qaic_config=qaic_config, moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) self.vision_model.onnx_path = vision_wrapper.vision_model.onnx_path @@ -1979,6 +2021,7 @@ def compile( qaic_config=qaic_config, layerwise_window_size=layerwise_window_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) @@ -2071,6 +2114,7 @@ def compile( _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, offload_pt_weights=offload_pt_weights, + use_weight_free_export=use_weight_free_export, ) if layerwise_cache_probe: return self.lang_model.onnx_path @@ -2117,7 +2161,7 @@ def compile( compiler_options["node_precision_info"] = self.model.get_npi_file(self.model.name_or_path) if not skip_lang: - custom_io_lang = {} + custom_io_lang = {"vision_embeds": CUSTOM_IO_DTYPE_MAP[target_dtype]} for output_name in output_names["lang"]: if output_name.endswith("_RetainedState"): dtype = ( @@ -2146,7 +2190,9 @@ def compile( else: specializations = lang_specs[:1] qpc_key = "lang_prefill_qpc_path" - elif prefill_seq_len == 1: + elif prefill_seq_len == 1 and not ( + self.continuous_batching and full_batch_size is not None and full_batch_size != batch_size + ): if self.comp_ctx_lengths_decode is not None: specializations = lang_specs[-len(self.comp_ctx_lengths_decode) :] else: @@ -2387,8 +2433,21 @@ def kv_offload_generate( } } - vision_inputs_fp16 = {"pixel_values", "image_masks"} - vision_inputs.update({k: vision_inputs[k].astype("float16") for k in vision_inputs_fp16 if k in vision_inputs}) + for input_name in ("pixel_values", "image_masks"): + if input_name not in vision_inputs or input_name not in vision_session.binding_index_map: + continue + binding = vision_session.bindings[vision_session.binding_index_map[input_name]] + vision_inputs[input_name] = vision_inputs[input_name].astype( + vision_session.aic_to_np_dtype_mapping[binding.type] + ) + + # Required for KIMI-K25 + grid_thws_val = inputs.pop("grid_thws", None) + if grid_thws_val is not None: + h_val = int(grid_thws_val[0, 1].item()) + w_val = int(grid_thws_val[0, 2].item()) + vision_inputs["h_shape"] = np.ones((h_val), dtype=np.int64) + vision_inputs["w_shape"] = np.ones((w_val), dtype=np.int64) vision_start = perf_counter() @@ -2500,6 +2559,9 @@ def kv_offload_generate( if self._write_io_dir is not None: write_io_files(lang_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + if "image_idx_output" in outputs: + lang_inputs["image_idx"] = chunk_inputs["image_idx"] + prefill_time = perf_counter() - lang_start + vision_end - vision_start # Skip inputs/outputs again lang_session.skip_buffers( @@ -3423,6 +3485,7 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel): AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, + PackQuantizedInt4ToMatMulNBitsTransform, Mxfp4GptOssExpertDequantizeTransform, CustomOpsTransform, KVCacheTransform, @@ -3431,6 +3494,12 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel): ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def prefill( self, @@ -3706,7 +3775,9 @@ def get_seq_len_and_handle_specialized_prefill_model( self.hash_params["moe_prefill_num_packed_chunks"] = num_packed_chunks if self.model.config.model_type in {"qwen3_moe", "gpt_oss", "glm4_moe"}: return max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) - return constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + seq_len = max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) + self.hash_params["chunking_seq_len"] = seq_len + return seq_len num_q_blocks = ( self.hash_params["blocking_config"].num_q_blocks if self.hash_params.get("blocking_kwargs", None) else None @@ -3872,13 +3943,13 @@ def export( # TODO: move this to a DA Serving utility class 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 []): + 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!" + ) seq_len = self.get_seq_len_and_handle_specialized_prefill_model( prefill_seq_len=prefill_seq_len, enable_chunking=enable_chunking, @@ -3889,6 +3960,10 @@ def export( kv_cache_shape[2] = ( seq_len + (sliding_window if sliding_window is not None else 0) if enable_chunking else seq_len ) + else: + seq_len = max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) + self.hash_params["chunking_seq_len"] = seq_len + else: self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) self.hash_params.pop("prefill_only", None) @@ -3999,8 +4074,14 @@ def export( example_inputs["compressed_kvs"][i].append( torch.zeros(pkv_cache[0][1].shape, dtype=self.model.config.torch_dtype) ) - dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"} - dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"} + dynamic_axes[f"compressed_kv.{i}"] = { + 0: "full_batch_size" if self.continuous_batching else "batch_size", + 2: "ctx_len", + } + dynamic_axes[f"k_pe.{i}"] = { + 0: "full_batch_size" if self.continuous_batching else "batch_size", + 2: "ctx_len", + } output_names.append(f"compressed_kv.{i}_RetainedState") output_names.append(f"k_pe.{i}_RetainedState") else: @@ -4088,6 +4169,17 @@ def _legacyify_cache(obj): output_names = apply_kv_cache_prefix(output_names, kv_cache_prefix) self.hash_params["kv_cache_prefix"] = kv_cache_prefix + if prefill_only: + effective_prefill_seq_len = ( + prefill_seq_len if prefill_seq_len is not None else seq_len if enable_chunking else None + ) + assert effective_prefill_seq_len is not None, "prefill_seq_len must be provided when prefill_only is True" + num_q_blocks_ffn = effective_prefill_seq_len // constants.EXPERT_BLOCKING_PACKED_CHUNK_SIZE + num_q_blocks_ffn = num_q_blocks_ffn if num_q_blocks_ffn > 0 else 1 + model_body = getattr(self.model, "model", None) + if model_body is not None: + setattr(model_body, "num_q_blocks_ffn", num_q_blocks_ffn) + if QEFFBaseModel._layerwise_active: return self._export_layerwise( example_inputs, diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 57c85f2fec..8ba55b67fe 100755 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -324,6 +324,7 @@ QEffDisentangledSelfAttention, ) from QEfficient.transformers.models.deepseek_v3.modeling_deepseek import ( + QEffDeepseekMoEGate, QEffDeepseekV3Attention, QEffDeepseekV3CustomRMSNormAIC, QEffDeepseekV3DecoderLayer, @@ -442,6 +443,11 @@ QEffInternVisionEmbeddings, QEffInternVLModel, ) +from QEfficient.transformers.models.kimi_k25.modeling_kimi_k25 import ( + QEffKimiK25ForConditionalGeneration, + QEffLearnable2DInterpPosEmbDivided_fixed, + QEffMoonViT3dEncoder, +) from QEfficient.transformers.models.llama.modeling_llama import ( QEffLlamaAttention, QEffLlamaDecoderLayer, @@ -1315,6 +1321,21 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "RMSNorm": { "forward": QEFFGrok1CustomRMSNormAIC.forward, }, + "KimiK25ForConditionalGeneration": { + "_qeff_merge_input_ids_with_image_features": QEffKimiK25ForConditionalGeneration._qeff_merge_input_ids_with_image_features, + "get_qeff_vision_encoder": QEffKimiK25ForConditionalGeneration.get_qeff_vision_encoder, + "get_qeff_language_decoder": QEffKimiK25ForConditionalGeneration.get_qeff_language_decoder, + "get_specializations": QEffKimiK25ForConditionalGeneration.get_specializations, + "get_onnx_dynamic_axes": QEffKimiK25ForConditionalGeneration.get_onnx_dynamic_axes, + "get_output_names": QEffKimiK25ForConditionalGeneration.get_output_names, + "get_dummy_inputs": QEffKimiK25ForConditionalGeneration.get_dummy_inputs, + }, + "MoonViT3dEncoder": { + "__qeff_init__": QEffMoonViT3dEncoder.__qeff_init__, + }, + "Learnable2DInterpPosEmbDivided_fixed": { + "__qeff_init__": QEffLearnable2DInterpPosEmbDivided_fixed.__qeff_init__, + }, "DeepseekV3ForCausalLM": { "forward": QEffDeepseekV3ForCausalLM.forward, "get_submodules_for_export": QEffDeepseekV3ForCausalLM.get_submodules_for_export, @@ -1326,8 +1347,11 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): }, "DeepseekV3MoE": { "forward": QEffDeepseekV3MoE.forward, - "moe": QEffDeepseekV3MoE.moe, + "moe_weights_as_activations": QEffDeepseekV3MoE.moe_weights_as_activations, + "moe_waa_unpack": QEffDeepseekV3MoE.moe_waa_unpack, + "original_moe": QEffDeepseekV3MoE.original_moe, "__qeff_init__": QEffDeepseekV3MoE.__qeff_init__, + "gate.forward": QEffDeepseekMoEGate.forward, }, "DeepseekV3Attention": { "forward": QEffDeepseekV3Attention.forward, @@ -1350,8 +1374,9 @@ class PrefillOnlyExternalModuleMapperTransform(ExternalModuleMapperTransform): _match_string_replace_method = { "DeepseekV3MoE": { "forward": QEffPrefillOnlyDeepseekV3MoE.forward, - "moe": QEffPrefillOnlyDeepseekV3MoE.moe, "__qeff_init__": QEffPrefillOnlyDeepseekV3MoE.__qeff_init__, + "_forward_expert_blocked": QEffPrefillOnlyDeepseekV3MoE._forward_expert_blocked, + "_cumsum_scatter_gather_update_expert_blocked": QEffPrefillOnlyDeepseekV3MoE._cumsum_scatter_gather_update_expert_blocked, }, } @@ -1361,7 +1386,7 @@ class RevertPrefillOnlyExternalModuleMapperTransform(ExternalModuleMapperTransfo _match_string_replace_method = { "DeepseekV3MoE": { "forward": QEffDeepseekV3MoE.forward, - "moe": QEffDeepseekV3MoE.moe, + "moe": QEffDeepseekV3MoE.moe_waa_unpack, "__qeff_init__": QEffDeepseekV3MoE.__qeff_init__, }, } diff --git a/QEfficient/transformers/quantizers/quant_transforms.py b/QEfficient/transformers/quantizers/quant_transforms.py index c3f204e629..8c2b997868 100644 --- a/QEfficient/transformers/quantizers/quant_transforms.py +++ b/QEfficient/transformers/quantizers/quant_transforms.py @@ -6,6 +6,8 @@ # ----------------------------------------------------------------------------- import torch +from compressed_tensors.compressors import PackedQuantizationCompressor +from compressed_tensors.linear.compressed_linear import CompressedLinear from torch import nn from transformers import AutoConfig from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts @@ -69,6 +71,84 @@ def mutate(cls, original_module: nn.Module, parent_module: nn.Module): return new_module +class PackQuantizedInt4ToMatMulNBitsTransform(ModuleMutatorTransform): + """ + This transform is used to pack the quantized int4 weights into a format that can be used by the MatMulNBits kernel. + It is used for the ONNX export of the quantized model. + """ + + _match_class = CompressedLinear + + @classmethod + def _is_pack_quantized_int4_linear(cls, module): + quantization_scheme = getattr(module, "quantization_scheme", None) + quantization_args = getattr(quantization_scheme, "weights", None) + return ( + isinstance(module, nn.Linear) + and hasattr(module, "weight_packed") + and hasattr(module, "weight_scale") + and quantization_args is not None + and getattr(quantization_args, "num_bits", None) == 4 + and getattr(quantization_args, "type", None) == "int" + and getattr(quantization_args, "strategy", None) == "group" + ) + + @classmethod + def _matches(cls, module): + return isinstance(module, cls._match_class) or cls._is_pack_quantized_int4_linear(module) + + @classmethod + def apply(cls, model: nn.Module): + transformed = False + for name, module in model.named_children(): + if cls._matches(module): + setattr(model, name, cls.mutate(module, model)) + transformed = True + else: + _, child_transformed = cls.apply(module) + transformed = transformed or child_transformed + + if cls._matches(model): + model = cls.mutate(model, None) + transformed = True + + return model, transformed + + @classmethod + def mutate(cls, original_module, parent_module): + # add compressor.decompress to get the decompressed weight + # and then package into matmulnbit + if isinstance(original_module, CompressedLinear): + assert isinstance(original_module.compressor, PackedQuantizationCompressor), ( + f"Only {PackedQuantizationCompressor} supported for now" + ) + fp_weight = original_module.compressor.decompress_module(original_module) + else: + from compressed_tensors.compressors.base import decompress_module as ct_decompress_module + + ct_decompress_module(original_module) + fp_weight = original_module.weight + scales = original_module.weight_scale + # assuming symmetric quantization + quantization_args = original_module.quantization_scheme.weights + zeros = (torch.zeros_like(scales) + pow(2, (quantization_args.num_bits - 1))).to(torch.uint8) + g_idx = torch.arange(original_module.in_features // quantization_args.group_size).repeat_interleave( + quantization_args.group_size + ) + original_module.weight = torch.nn.Parameter(fp_weight) + assert quantization_args.type == "int", "uint is not tested yet" + new_module = QuantLinearORT( + quantization_args.num_bits, + quantization_args.group_size, + original_module.in_features, + original_module.out_features, + original_module.bias is not None, + ) + new_module.bias = original_module.bias if original_module.bias is not None else None + new_module.pack(original_module, scales, zeros, g_idx) + return new_module + + class GPTQToMatmulNbitsTransform(ModuleMutatorTransform): """ A transformation class that mutates a ``QuantLinearGPTQ`` module to a ``QuantLinearORT`` diff --git a/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py b/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py index 81c6b81cdf..43ffb32ee2 100644 --- a/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py +++ b/QEfficient/transformers/quantizers/quantizer_compressed_tensors.py @@ -12,6 +12,7 @@ import torch from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts from transformers.quantizers.quantizer_compressed_tensors import CompressedTensorsHfQuantizer +from transformers.utils import is_compressed_tensors_available from transformers.utils.quantization_config import CompressedTensorsConfig, QuantizationConfigMixin, QuantizationMethod from QEfficient.transformers.quantizers.quantizer_utils import blockwise_dequantize, get_keys_to_not_convert @@ -390,7 +391,55 @@ def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str class QEffCompressedTensorsConfig(CompressedTensorsConfig): - def __init__( + def handle_pack_quantized_init( + self, + config_groups=None, + format="dense", + quantization_status="initialized", + kv_cache_scheme=None, + global_compression_ratio=None, + ignore=None, + sparsity_config=None, + quant_method="compressed-tensors", + run_compressed: bool = True, + **kwargs, + ): + if is_compressed_tensors_available(): + from compressed_tensors.config import SparsityCompressionConfig + from compressed_tensors.quantization import QuantizationConfig + else: + raise ImportError( + "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`." + ) + self.quantization_config = None + self.sparsity_config = None + + self.run_compressed = run_compressed + assert self.run_compressed, "pack-quantized needs to have run_compressed set to True" + + # parse from dict to load nested QuantizationScheme objects + if config_groups or kv_cache_scheme: + self.quantization_config = QuantizationConfig.model_validate( + { + "config_groups": config_groups, + "quant_method": quant_method, + "format": format, + "quantization_status": quantization_status, + "kv_cache_scheme": kv_cache_scheme, + "global_compression_ratio": global_compression_ratio, + "ignore": ignore, + **kwargs, + } + ) + + if sparsity_config: + self.sparsity_config = SparsityCompressionConfig.load_from_registry( + sparsity_config.get("format"), **sparsity_config + ) + + self.quant_method = QuantizationMethod.COMPRESSED_TENSORS + + def handle_fp8_init( self, config_groups=None, format="dense", @@ -480,7 +529,50 @@ def __init__( self.quant_method = QuantizationMethod.COMPRESSED_TENSORS + def __init__( + self, + config_groups=None, + format="dense", + quantization_status="initialized", + kv_cache_scheme=None, + global_compression_ratio=None, + ignore=None, + sparsity_config=None, + quant_method="compressed-tensors", + run_compressed: bool = None, + **kwargs, + ): + if format == "pack-quantized": + self.handle_pack_quantized_init( + config_groups=config_groups, + format=format, + quantization_status=quantization_status, + kv_cache_scheme=kv_cache_scheme, + global_compression_ratio=global_compression_ratio, + ignore=ignore, + sparsity_config=sparsity_config, + quant_method=quant_method, + run_compressed=True if run_compressed is None else run_compressed, + **kwargs, + ) + else: + self.handle_fp8_init( + config_groups=config_groups, + format=format, + quantization_status=quantization_status, + kv_cache_scheme=kv_cache_scheme, + global_compression_ratio=global_compression_ratio, + ignore=ignore, + sparsity_config=sparsity_config, + quant_method=quant_method, + run_compressed=False if run_compressed is None else run_compressed, + **kwargs, + ) + def to_dict(self): + if getattr(self.quantization_config, "format", None) == "pack-quantized": + return super().to_dict() + return { "quantization_config": { "config_groups": self.config_groups, @@ -501,39 +593,58 @@ def to_dict(self): class QEffCompressedTensorsFP8Quantizer(CompressedTensorsHfQuantizer): requires_calibration = False - def __init__(self, quantization_config, **kwargs): - # TODO: check if more checks are required - if not isinstance(quantization_config, QEffCompressedTensorsConfig): - raise TypeError( - f"Only {QEffCompressedTensorsConfig} is supported for initialization got {type(quantization_config)}" - ) - self.run_compressed = quantization_config.run_compressed - self.quantization_config = quantization_config - - # -- Handle extra kwargs below -- - self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) - self.modules_to_not_convert = list( - set(self.modules_to_not_convert if self.modules_to_not_convert else []) - | set(self.quantization_config.ignore if self.quantization_config.ignore else []) + @staticmethod + def is_pack_quantized(quant_config): + return ( + hasattr(quant_config, "quantization_config") + and hasattr(quant_config.quantization_config, "format") + and quant_config.quantization_config.format == "pack-quantized" ) - self.pre_quantized = kwargs.pop("pre_quantized", True) - if not self.pre_quantized and self.requires_calibration: - raise ValueError( - f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." - f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " - f"pass `pre_quantized=True` while knowing what you are doing." + def __init__(self, quantization_config, **kwargs): + if self.is_pack_quantized(quantization_config): + super().__init__(quantization_config, **kwargs) + else: + if not isinstance(quantization_config, QEffCompressedTensorsConfig): + raise TypeError( + f"Only {QEffCompressedTensorsConfig} is supported for initialization got {type(quantization_config)}" + ) + self.run_compressed = quantization_config.run_compressed + self.quantization_config = quantization_config + + # -- Handle extra kwargs below -- + self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) + self.modules_to_not_convert = list( + set(self.modules_to_not_convert if self.modules_to_not_convert else []) + | set(self.quantization_config.ignore if self.quantization_config.ignore else []) ) + self.pre_quantized = kwargs.pop("pre_quantized", True) + + if not self.pre_quantized and self.requires_calibration: + raise ValueError( + f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." + f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " + f"pass `pre_quantized=True` while knowing what you are doing." + ) def validate_environment(self, *args, **kwargs): + if self.is_pack_quantized(self.quantization_config): + return super().validate_environment(*args, **kwargs) return True def update_torch_dtype(self, torch_dtype): + if self.is_pack_quantized(self.quantization_config): + return super().update_torch_dtype(torch_dtype) + if torch_dtype not in [None, torch.float32]: logger.warning(f"Requested dtype {torch_dtype} is not supported, overriding to float32") return torch.float32 def _process_model_before_weight_loading(self, model, **kwargs): + if self.is_pack_quantized(self.quantization_config): + super()._process_model_before_weight_loading(model, **kwargs) + return + if self.quantization_config.targets != ["Linear"]: raise NotImplementedError( f"Only Linear layer with FP8 quantization are supported got targets = {self.quantization_config.targets}" @@ -561,9 +672,14 @@ def replace_linear_with_fp8_dequant_layer(module): replace_linear_with_fp8_dequant_layer(model) def _process_model_after_weight_loading(self, model, **kwargs): + if self.is_pack_quantized(self.quantization_config): + super()._process_model_after_weight_loading(model, **kwargs) + return pass def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: + if self.is_pack_quantized(self.quantization_config): + return super().update_missing_keys_after_loading(model, missing_keys=missing_keys, prefix=prefix) return missing_keys def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str = None) -> List[str]: diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index c8a64adefb..f26b09087b 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -148,6 +148,7 @@ def get_default_aic_hw_version() -> str: DEFAULT_AIC_HW_VERSION = get_default_aic_hw_version() ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL = 100 +EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) # Generic config key aliases used across model families. ATTENTION_HEAD_CONFIG_KEYS = ("num_attention_heads", "n_head", "n_heads", "num_heads") KV_HEAD_CONFIG_KEYS = ("num_key_value_heads", "n_kv_heads", "num_kv_heads", "effective_n_kv_heads") @@ -185,7 +186,11 @@ def get_default_aic_hw_version() -> str: LLAMA4_ATTENTION_CHUNK_SIZE = 8192 LLAMA4_MAX_POSITION_EMBEDDINGS = 65536 -# DeepSeek Kimi-k2 Constant +# DeepSeek Kimi-k2.5 Constants +KIMI_PATCH_SIZE = 14 +KIMI_EXAMPLE_IMAGE_NUM_IMAGE_TOKENS = 600 +KIMI_EXAMPLE_IMAGE_NUM_PATCHES_HEIGHT = 30 +KIMI_EXAMPLE_IMAGE_NUM_PATCHES_WIDTH = 80 MAX_POSITION_EMBEDDINGS = 32768 FP16_BYTES = 2 DEFAULT_NUM_HEADS = 64 diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index bc61f868a2..9218191e2a 100755 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -69,6 +69,7 @@ def convert_dynamic_axes_to_dynamic_shapes( torch.onnx.export(dynamic_shapes=...). """ max_seq_len = getattr(model_config, "max_position_embeddings", 1024) + max_image_dim = max(max_seq_len, 65536) model_type = getattr(model_config, "model_type", None) batch_min = 1 if model_type == "gpt_oss" else 2 @@ -87,6 +88,12 @@ def resolve_dim(dim_name: str): dim_registry[dim_name] = Dim("comp_ctx_lengths", min=DYNAMO_DIM_MIN_COMP_CTX_LENGTHS, max=max_seq_len) elif "ctx_len" in dim_name: dim_registry[dim_name] = Dim("ctx_len", min=2, max=max_seq_len) + elif dim_name == "num_patches": + dim_registry[dim_name] = Dim("num_patches", min=1, max=max_image_dim) + elif dim_name == "num_image_tokens": + dim_registry[dim_name] = Dim("num_image_tokens", min=1, max=max_image_dim) + elif dim_name in {"grid_h", "grid_w"}: + dim_registry[dim_name] = Dim(dim_name, min=1, max=max_image_dim) elif "sliding_window" in dim_name: dim_registry[dim_name] = Dim( "sliding_window", diff --git a/QEfficient/utils/test_utils.py b/QEfficient/utils/test_utils.py index 89dfb66029..52a692b965 100644 --- a/QEfficient/utils/test_utils.py +++ b/QEfficient/utils/test_utils.py @@ -98,6 +98,8 @@ def set_num_layers_vlm(config: AutoConfig, n_layer: int = -1): elif hasattr(config, "text_config"): config.text_config.num_hidden_layers = n_layer config.vision_config.num_hidden_layers = n_layer + if hasattr(config.vision_config, "vt_num_hidden_layers"): + config.vision_config.vt_num_hidden_layers = n_layer if hasattr(config.vision_config, "depth"): config.vision_config.depth = n_layer if hasattr(config.vision_config, "deepstack_visual_indexes"): @@ -487,6 +489,7 @@ class ModelConfig: } DUAL_QPC_MODELS = { + "moonshotai/Kimi-K2.5", "OpenGVLab/InternVL2_5-1B", "OpenGVLab/InternVL3_5-1B", "Qwen/Qwen2.5-VL-3B-Instruct", diff --git a/QEfficient/utils/torch_patches.py b/QEfficient/utils/torch_patches.py index 18cab51231..8e705eedf2 100644 --- a/QEfficient/utils/torch_patches.py +++ b/QEfficient/utils/torch_patches.py @@ -26,6 +26,7 @@ - _translate_fx_graph / _convert_fx_arg_to_onnx_arg nested tensor constants """ +import inspect from contextlib import contextmanager import torch @@ -337,3 +338,50 @@ def temporarily_enable_nested_compile_regions(model, target_classes=None): delattr(module, "forward") else: setattr(module, "forward", previous_forward) + + +@contextmanager +def temporarily_disable_nested_compile_regions(model, target_classes=None): + """ + Replace nested_compile_region-wrapped ``forward`` methods with their original + underlying functions for the duration of plain dynamo export (flat graph path). + + Used when use_weight_free_export=True and use_onnx_subfunctions=False so that + @nested_compile_region boundaries on decoder layer forward() methods do not + create unwanted subgraph splits during tracing. + """ + target_classes = tuple(target_classes) if target_classes else None + patched_modules = [] + + try: + for module in model.modules(): + if target_classes and not isinstance(module, target_classes): + continue + + bound_forward = getattr(module, "forward", None) + if bound_forward is None: + continue + + wrapped_forward = getattr(bound_forward, "__func__", bound_forward) + if getattr(wrapped_forward, "__qualname__", "") != "mark_compile_region..wrap..inner": + continue + + closure = getattr(wrapped_forward, "__closure__", None) or () + original_forward = next( + (cell.cell_contents for cell in closure if inspect.isfunction(cell.cell_contents)), + None, + ) + if original_forward is None: + continue + + previous_forward = module.__dict__.get("forward", _MISSING_INSTANCE_ATTR) + setattr(module, "forward", original_forward.__get__(module, type(module))) + patched_modules.append((module, previous_forward)) + + yield + finally: + for module, previous_forward in reversed(patched_modules): + if previous_forward is _MISSING_INSTANCE_ATTR: + delattr(module, "forward") + else: + setattr(module, "forward", previous_forward) diff --git a/README.md b/README.md index b7d8fe4b42..81b6bee2b5 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ *Latest news* :fire:
+- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [07/2026] Added `dynamo` flag to `QEFFAutoModelForCausalLM.export()` to support `torch.onnx.export` dynamo-based ONNX export for CausalLM models - [06/2026] Added support for Gemma4 models, [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) - [06/2026] Added support for Qwen3.6 model [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) diff --git a/docs/source/introduction.md b/docs/source/introduction.md index a02bbef8c7..616d5bbd84 100644 --- a/docs/source/introduction.md +++ b/docs/source/introduction.md @@ -22,6 +22,7 @@ For other models, there is comprehensive documentation to inspire upon the chang *Latest news* :
+- [07/2026] Added support for Kimi-K2.5 vision-language model [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) via `QEFFAutoModelForImageTextToText`. - [07/2026] Added `dynamo` flag to `QEFFAutoModelForCausalLM.export()` to support `torch.onnx.export` dynamo-based ONNX export for CausalLM models - [06/2026] Added support for Gemma4 models, [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) - [06/2026] Added support for Qwen3.6 model [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) diff --git a/docs/source/release_docs.md b/docs/source/release_docs.md index 2724ab6dd6..09d4e32931 100644 --- a/docs/source/release_docs.md +++ b/docs/source/release_docs.md @@ -227,6 +227,12 @@ Welcome to the official release of **Efficient Transformer Library v1.21.0**! Th - Executable via [`QEFFAutoModelForImageTextToText`](#QEFFAutoModelForImageTextToText) - [Mistral-3.1 Example Script](https://github.com/quic/efficient-transformers/blob/main/examples/image_text_to_text/models/mistral_vision/mistral3_example.py) +- **Kimi-K2.5 (Vision Language)** + - Executable via [`QEFFAutoModelForImageTextToText`](#QEFFAutoModelForImageTextToText) + - Supports `moonshotai/Kimi-K2.5` with `KimiK25ForConditionalGeneration` + - Supports dual-QPC (`kv_offload=True`) vision-language compilation + - [Kimi-K2.5 Vision Example Script](https://github.com/quic/efficient-transformers/blob/main/examples/kimi_k2/export_kimi_k25_vision.py) + - **Disaggregated serving ready via vLLM GPT-OSS** > **Note**: If running GPT-OSS models natively via vLLM, PR-685 of the qefficient library is required for Python 3.12 compatibility. @@ -279,7 +285,7 @@ Welcome to the official release of **Efficient Transformer Library v1.21.0**! Th - **Compute-Context-Length (CCL) support**: To optimize the throughput when handling very large context lengths - **Prefill/Decode Separation**: Support for GPT OSS using disaggregate serving models - **Continuous Batching (VLMs)**: Extended to Vision Language Models with multi-image handling - - Supported models: Llava, Llava_Next, Gemma3, Mistral3, InternVL2_5, InternVL3_5, Molmo + - Supported models: Llava, Llava_Next, Gemma3, Mistral3, InternVL2_5, InternVL3_5, Molmo, Kimi-K2.5 - **ONNX Sub-Functions**: Feature enabling more efficient model compilation and execution on hardware. Users can enable the feature by passing `use_onnx_subfunctions=True` during export - **Memory Profiling**: Built-in utilities for optimization analysis - **Extend on-device Sampling**: Extend on-device sampling to dual QPC VLMs and Guided decoding for on-device sampling diff --git a/docs/source/validate.md b/docs/source/validate.md index 48a0e5b0f8..be5009aac7 100644 --- a/docs/source/validate.md +++ b/docs/source/validate.md @@ -90,6 +90,7 @@ | **Qwen3_5MoeForConditionalGeneration** | Qwen3.5 | [Qwen/Qwen3.5-122B-A10B](https://huggingface.co/Qwen/Qwen3.5-122B-A10B)
[Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5ForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) | ✕ | ✔️ | ✕ | ✔️ | | **Qwen3_5MoeForConditionalGeneration** | Qwen3.6 | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) | ✕ | ✔️ | ✕ | ✔️ | +| **KimiK25ForConditionalGeneration** | Kimi-K2.5 | [moonshotai/Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) | ✕ | ✔️ | ✕ | ✔️ | | **Mistral3ForConditionalGeneration** | Mistral3| [mistralai/Mistral-Small-3.1-24B-Instruct-2503](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503)| ✕ | ✔️ | ✕ | ✕ | ### Vision-Language Reranker Models (Text + Image Scoring) diff --git a/examples/image_text_to_text/README.md b/examples/image_text_to_text/README.md index 612f9d400a..e9504b6b0c 100644 --- a/examples/image_text_to_text/README.md +++ b/examples/image_text_to_text/README.md @@ -91,6 +91,7 @@ Popular model families include: - InternVL - Molmo - LLaVA +- Kimi-K2.5 ### Model-Specific Examples @@ -107,6 +108,23 @@ Some models have specialized examples demonstrating advanced features: | **Granite** | [models/granite_vision/](models/granite_vision/) | | **InternVL** | [models/internvl/](models/internvl/) | | **Molmo** | [models/molmo/](models/molmo/) | +| **Kimi-K2.5** | [../kimi_k2/export_kimi_k25_vision.py](../kimi_k2/export_kimi_k25_vision.py) | + +### Kimi-K2.5 Vision + +Kimi-K2.5 uses the `moonshotai/Kimi-K2.5` remote-code model with the `KimiK25ForConditionalGeneration` architecture. The example script loads either the full checkpoint or a smaller layer subset, wraps it with `QEFFAutoModelForImageTextToText`, and compiles the vision-language path with Kimi-specific image-size specializations. + +```bash + +python ../kimi_k2/export_kimi_k25_vision.py \ + --model-path "$HF_HUB_CACHE/models--moonshotai--Kimi-K2.5/snapshots/" \ + --image-url "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" \ + --image-height 512 \ + --image-width 910 \ + --prompt "Describe this image." +``` + +For custom image sizes, pass only the input image pixel dimensions with `--image-height` and `--image-width`. Kimi-K2.5 compile derives the internal grid height, grid width, patch count, and image-token count from those pixel dimensions. For reranker examples, see [../reranker/](../reranker/). diff --git a/examples/kimi_k2/export_kimi_k25_dynamo.py b/examples/kimi_k2/export_kimi_k25_dynamo.py new file mode 100644 index 0000000000..f2b03bb1f1 --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_dynamo.py @@ -0,0 +1,366 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +"""Export, compile, and generate with a reduced Kimi K2.5 model slice. + +The ONNX export step uses the Dynamo path directly on the dual-QPC component +wrappers. The exported ONNX files are then passed to the normal QEfficient +compile and generation APIs. +""" + +import argparse +import importlib.util +import os +from io import BytesIO +from pathlib import Path + +import requests +import torch +from PIL import Image + +from QEfficient import QEFFAutoModelForImageTextToText + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +DEFAULT_EXPORT_DIR = Path.home() / ".cache" / "qeff" / "kimi_k25_dynamo_export" +DEFAULT_COMPILE_DIR = Path.home() / ".cache" / "qeff" / "kimi_k25_dynamo_compile" +DEFAULT_IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Export Kimi K2.5 ONNX with Dynamo, compile QPCs, and run generation.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--model-path", + type=Path, + default="/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611", + help="Local Kimi-K2.5 snapshot path. Uses the local HF cache/default snapshot when omitted.", + ) + parser.add_argument("--export-dir", type=Path, default=DEFAULT_EXPORT_DIR) + parser.add_argument("--compile-dir", type=Path, default=DEFAULT_COMPILE_DIR) + parser.add_argument("--vision-onnx-path", type=Path, default=None) + parser.add_argument("--lang-onnx-path", type=Path, default=None) + parser.add_argument("--vision-qpc-path", type=Path, default=None) + parser.add_argument("--lang-qpc-path", type=Path, default=None) + parser.add_argument("--component", choices=("lang", "vision", "both"), default="both") + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument("--prefill-seq-len", type=int, default=2) + parser.add_argument("--ctx-len", type=int, default=1024) + parser.add_argument("--num-devices", type=int, default=1) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--image-url", type=str, default=DEFAULT_IMAGE_URL) + parser.add_argument("--image-path", type=Path, default=None) + parser.add_argument("--image-height", type=int, default=None) + parser.add_argument("--image-width", type=int, default=None) + parser.add_argument("--generation-len", type=int, default=10) + parser.add_argument( + "--device-ids", + type=lambda device_ids: [int(device_id) for device_id in device_ids.strip("[]").split(",")], + default=[0], + help="Device IDs for generation, e.g. [0] or [0,1].", + ) + parser.add_argument("--mxfp6-matmul", action="store_true") + parser.add_argument("--mxint8-kv-cache", action="store_true") + parser.add_argument("--mos", type=int, default=1) + parser.add_argument("--aic-enable-depth-first", action="store_true") + parser.add_argument("--skip-export", action="store_true") + parser.add_argument("--skip-compile", action="store_true") + parser.add_argument("--skip-generate", action="store_true") + parser.add_argument( + "--use-onnx-subfunctions", + action="store_true", + help="Enable repeated-subgraph ONNX functions during Dynamo export.", + ) + parser.add_argument( + "--keep-weights", + action="store_true", + help="Keep PyTorch weights resident after export instead of offloading them to meta tensors.", + ) + parser.add_argument( + "--use-weight-free-export", + action="store_true", + help="Export ONNX graphs without embedded weights and emit weight_spec.json metadata.", + ) + args = parser.parse_args() + if (args.image_height is None) != (args.image_width is None): + parser.error("--image-height and --image-width must be provided together.") + if not args.skip_generate and args.component != "both": + parser.error("Generation requires --component both so both vision and language QPCs are available.") + if args.skip_export and (args.vision_onnx_path is None or args.lang_onnx_path is None) and not args.skip_compile: + parser.error("--skip-export requires --vision-onnx-path and --lang-onnx-path unless --skip-compile is set.") + if args.skip_compile and not args.skip_generate and (args.vision_qpc_path is None or args.lang_qpc_path is None): + parser.error("--skip-compile generation requires --vision-qpc-path and --lang-qpc-path.") + return args + + +def validate_dynamo_torch_version(): + torch_version = torch.__version__ + major, minor = (int(part) for part in torch_version.split("+")[0].split(".")[:2]) + if (major, minor) < (2, 13): + raise RuntimeError( + f"Kimi K2.5 Dynamo export requires PyTorch >= 2.13, but {torch_version} is installed. " + "Install the Dynamo requirements before running this example." + ) + + +def configure_qaic_tool_path(): + qaic_exec_path = Path("/opt/qti-aic/exec") + if qaic_exec_path.exists(): + os.environ["PATH"] = f"{qaic_exec_path}{os.pathsep}{os.environ.get('PATH', '')}" + + +def build_qeff_model(args): + model_path = args.model_path + subset_model_path = None + model_ref = model_path + if args.use_weight_free_export: + subset_model_path = args.export_dir.expanduser().resolve() / "kimi_k25_weightfree_subset" + model_ref = subset_model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + subset_model_path=subset_model_path, + ) + model.eval() + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText( + model, qaic_config=qaic_config, pretrained_model_name_or_path=str(model_ref) + ) + qeff_model.model.eval() + qeff_model.vision_model.model.eval() + qeff_model.lang_model.model.eval() + qeff_model.transform( + ctx_len=args.ctx_len, + seq_len=args.prefill_seq_len, + bs=1, + num_devices=args.num_devices, + qaic_config=qaic_config, + aic_num_cores=args.num_cores, + ) + precompute_vision_rope_cache(qeff_model) + return qeff_model, tokenizer, processor, qaic_config + + +def precompute_vision_rope_cache(qeff_model): + rope_2d = qeff_model.vision_model.model.model.vision_tower.encoder.rope_2d + rope_2d._ensure_precomputed_freqs(torch.device("cpu")) + + +def get_component_export_args(qeff_model, prefill_seq_len: int): + inputs = qeff_model.model.get_dummy_inputs( + kv_offload=True, + continuous_batching=qeff_model.continuous_batching, + prefill_seq_len=prefill_seq_len, + ) + normalize_nested_cache_inputs_for_dynamo(inputs) + dynamic_axes = qeff_model.model.get_onnx_dynamic_axes( + kv_offload=True, + continuous_batching=qeff_model.continuous_batching, + comp_ctx_lengths=qeff_model.comp_ctx_lengths_decode, + ) + output_names = qeff_model.model.get_output_names(kv_offload=True) + add_static_dynamic_axes_for_dynamo(inputs, dynamic_axes) + remove_vision_output_dynamic_axes_for_dynamo(dynamic_axes) + freeze_batch_axes_for_dynamo(dynamic_axes) + return inputs, output_names, dynamic_axes + + +def normalize_nested_cache_inputs_for_dynamo(inputs): + for component_inputs in inputs.values(): + if "compressed_kvs" in component_inputs: + component_inputs["compressed_kvs"] = [tuple(layer) for layer in component_inputs["compressed_kvs"]] + if "past_key_values" in component_inputs: + component_inputs["past_key_values"] = [list(layer) for layer in component_inputs["past_key_values"]] + + +def add_static_dynamic_axes_for_dynamo(inputs, dynamic_axes): + nested_cache_inputs = {"past_key_values", "compressed_kvs"} + for component_name, component_inputs in inputs.items(): + component_dynamic_axes = dynamic_axes[component_name] + for input_name in component_inputs: + if input_name not in nested_cache_inputs: + component_dynamic_axes.setdefault(input_name, {}) + + +def remove_vision_output_dynamic_axes_for_dynamo(dynamic_axes): + dynamic_axes["vision"].pop("vision_embeds", None) + + +def freeze_batch_axes_for_dynamo(dynamic_axes): + for component_dynamic_axes in dynamic_axes.values(): + for axes_map in component_dynamic_axes.values(): + for axis_idx, dim_name in tuple(axes_map.items()): + if dim_name in {"batch_size", "full_batch_size"}: + axes_map.pop(axis_idx) + + +def export_vision(qeff_model, inputs, output_names, dynamic_axes, args): + return qeff_model.vision_model._export( + inputs["vision"], + output_names=output_names["vision"], + dynamic_axes=dynamic_axes["vision"], + export_dir=args.export_dir, + offload_pt_weights=False, + dynamo=True, + use_onnx_subfunctions=args.use_onnx_subfunctions, + use_weight_free_export=args.use_weight_free_export, + ) + + +def export_language(qeff_model, inputs, output_names, dynamic_axes, args): + qeff_model.lang_model.hash_params["prefill_only"] = False + onnx_path = qeff_model.lang_model._export( + inputs["lang"], + output_names=output_names["lang"], + dynamic_axes=dynamic_axes["lang"], + export_dir=args.export_dir, + offload_pt_weights=not args.keep_weights, + dynamo=True, + use_onnx_subfunctions=args.use_onnx_subfunctions, + use_weight_free_export=args.use_weight_free_export, + ) + return onnx_path + + +def export_components(qeff_model, inputs, output_names, dynamic_axes, args): + exported_paths = {"vision": args.vision_onnx_path, "lang": args.lang_onnx_path} + if args.skip_export: + return {component_name: path for component_name, path in exported_paths.items() if path is not None} + if args.component in {"vision", "both"}: + exported_paths["vision"] = export_vision(qeff_model, inputs, output_names, dynamic_axes, args) + if args.component in {"lang", "both"}: + exported_paths["lang"] = export_language(qeff_model, inputs, output_names, dynamic_axes, args) + return exported_paths + + +def load_generation_image(args): + if args.image_path is not None: + image = Image.open(args.image_path).convert("RGB") + else: + response = requests.get(args.image_url, timeout=30) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) + return image + + +def compile_components(qeff_model, exported_paths, image, qaic_config, args): + if args.skip_compile: + qpc_paths = {"vision_qpc_path": str(args.vision_qpc_path), "lang_qpc_path": str(args.lang_qpc_path)} + qeff_model.vision_model.qpc_path = qpc_paths["vision_qpc_path"] + qeff_model.lang_model.qpc_path = qpc_paths["lang_qpc_path"] + qeff_model.qpc_paths = qpc_paths + return qpc_paths + + qpc_paths = qeff_model.compile( + vision_onnx_path=str(exported_paths.get("vision")), + lang_onnx_path=str(exported_paths.get("lang")), + compile_dir=str(args.compile_dir), + qaic_config=qaic_config, + prefill_seq_len=args.prefill_seq_len, + ctx_len=args.ctx_len, + num_cores=args.num_cores, + num_devices=args.num_devices, + mxfp6_matmul=args.mxfp6_matmul, + mxint8_kv_cache=args.mxint8_kv_cache, + skip_vision=args.component == "lang", + skip_lang=args.component == "vision", + aic_enable_depth_first=args.aic_enable_depth_first, + mos=args.mos, + image_height=image.height if image is not None else args.image_height, + image_width=image.width if image is not None else args.image_width, + use_weight_free_export=args.use_weight_free_export, + ) + return qpc_paths + + +def build_generation_inputs(processor, qeff_model, image, args): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + return inputs + + +def generate(qeff_model, tokenizer, processor, image, args): + inputs = build_generation_inputs(processor, qeff_model, image, args) + output = qeff_model.generate( + inputs=inputs, + device_ids=args.device_ids, + generation_len=args.generation_len, + image_height=image.height, + image_width=image.width, + ) + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + return output + + +def main(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + args = parse_args() + validate_dynamo_torch_version() + configure_qaic_tool_path() + qeff_model, tokenizer, processor, qaic_config = build_qeff_model(args) + inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + + exported_paths = export_components(qeff_model, inputs, output_names, dynamic_axes, args) + + for component_name, onnx_path in exported_paths.items(): + print(f"{component_name} ONNX exported to: {onnx_path}") + + image = load_generation_image(args) if not args.skip_generate else None + if not args.skip_compile or not args.skip_generate: + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + print(f"QPC paths: {qpc_paths}") + if not args.skip_generate: + generate(qeff_model, tokenizer, processor, image, args) + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py new file mode 100644 index 0000000000..93f46cd6b3 --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -0,0 +1,396 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import importlib.util +from io import BytesIO +from pathlib import Path +from time import perf_counter + +import numpy as np +import requests +import torch +from PIL import Image + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +load_kimi_k25_class = load_kimi_utils.load_kimi_k25_class +load_layer_subset_model = load_kimi_utils.load_layer_subset_model +parse_expert_ids = load_kimi_utils.parse_expert_ids +prepare_config = load_kimi_utils.prepare_config +set_deterministic = load_kimi_utils.set_deterministic + +PREFILL_SEQ_LEN = 512 +CTX_LEN = 2048 +BS = 1 +GENERATION_LEN = 10 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Run Kimi K2.5 vision disaggregated vision -> prefill -> decode flow.") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument( + "--full-model", + action="store_true", + help="Load the full model. By default, the script loads a small layer subset for faster startup.", + ) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument( + "--image-height", + type=int, + default=None, + help="Image height in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image height.", + ) + parser.add_argument( + "--image-width", + type=int, + default=None, + help="Image width in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image width.", + ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--prefill-seq-len", type=int, default=PREFILL_SEQ_LEN) + parser.add_argument("--ctx-len", type=int, default=CTX_LEN) + parser.add_argument("--generation-len", type=int, default=GENERATION_LEN) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument("--vision-num-devices", type=int, default=1) + parser.add_argument("--lang-num-devices", type=int, default=4) + parser.add_argument("--mxfp6-matmul", action="store_true") + parser.add_argument("--mxint8-kv-cache", action="store_true") + args = parser.parse_args() + if (args.image_height is None) != (args.image_width is None): + parser.error("--image-height and --image-width must be provided together.") + return args + + +def _clone_inputs(inputs): + return {key: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for key, value in inputs.items()} + + +def _numpy(value): + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _session_input_names(session: QAICInferenceSession) -> set[str]: + input_names = set(session.input_names) + input_names.update(name.rsplit("/", 1)[-1] for name in session.input_names) + return input_names + + +def _cast_for_session(session: QAICInferenceSession, name: str, value: np.ndarray) -> np.ndarray: + binding_index = session.binding_index_map.get(name) + if binding_index is None: + return value + dtype = session.aic_to_np_dtype_mapping[session.bindings[binding_index].type] + return value.astype(dtype, copy=False) + + +def _filter_session_inputs(session: QAICInferenceSession, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + input_names = _session_input_names(session) + return {name: _cast_for_session(session, name, value) for name, value in inputs.items() if name in input_names} + + +def _resolve_qpc_path(qpc_paths, key: str): + if isinstance(qpc_paths, dict): + qpc_path = qpc_paths.get(key) + if qpc_path is None: + raise KeyError(f"Missing {key!r} in compile output keys: {list(qpc_paths.keys())}") + return qpc_path + return qpc_paths + + +def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs: dict[str, np.ndarray]): + for output_name, value in source_outputs.items(): + output_basename = output_name.rsplit("/", 1)[-1] + if output_basename.endswith("_RetainedState"): + target_inputs[output_basename.removesuffix("_RetainedState")] = value + + +def _get_next_token_ids(logits: np.ndarray) -> np.ndarray: + logits = np.asarray(logits) + return logits[:, -1, :].argmax(axis=-1).astype(np.int64).reshape(BS, 1) + + +def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, args, image: Image.Image): + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + common_compile_kwargs = { + "qaic_config": qaic_config, + "batch_size": BS, + "ctx_len": args.ctx_len, + "image_height": image.height, + "image_width": image.width, + "num_cores": args.num_cores, + "mxfp6_matmul": args.mxfp6_matmul, + "mxint8_kv_cache": args.mxint8_kv_cache, + "split_model_io": True, + "mos": 1, + "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, + "layerwise": False, + } + + print("Compiling vision QPC...") + vision_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + skip_vision=False, + skip_lang=True, + num_devices=args.vision_num_devices, + **common_compile_kwargs, + ) + + print("Compiling prefill QPC...") + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + prefill_only=True, + enable_chunking=True, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + print("Compiling decode QPC...") + decode_qpc_path = qeff_model.compile( + prefill_seq_len=1, + prefill_only=False, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + return vision_qpc_path, prefill_qpc_path, decode_qpc_path + + +def _run_disagg_generation( + inputs: dict[str, torch.Tensor], + vision_session: QAICInferenceSession, + prefill_session: QAICInferenceSession, + decode_session: QAICInferenceSession, + *, + prefill_seq_len: int, + generation_len: int, +) -> np.ndarray: + inputs = {name: _numpy(value) for name, value in _clone_inputs(inputs).items()} + input_ids_length = inputs["input_ids"].shape[1] + num_chunks = -(input_ids_length // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + + inputs["input_ids"] = np.pad( + inputs["input_ids"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=1, + ) + inputs["attention_mask"] = np.pad( + inputs["attention_mask"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=0, + ) + + grid_thws = inputs.pop("grid_thws").astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + } + + vision_start = perf_counter() + print("Running vision QPC...") + vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) + vision_session.deactivate() + vision_time = perf_counter() - vision_start + + vision_embeds = vision_outputs.get("vision_embeds") + if vision_embeds is None: + raise RuntimeError(f"Vision QPC did not return vision_embeds. Outputs: {vision_outputs.keys()}") + + lang_inputs = { + "input_ids": inputs["input_ids"].astype(np.int64), + "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), + "vision_embeds": vision_embeds, + "image_idx": np.zeros((BS, 1), dtype=np.int64), + } + + prefill_start = perf_counter() + print("Running prefill QPC...") + prefill_session.set_buffers(vision_outputs) + chunk_inputs = lang_inputs.copy() + prefill_outputs = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = (chunk_idx + 1) * prefill_seq_len + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, start:end] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][:, start:end] + prefill_outputs = prefill_session.run(_filter_session_inputs(prefill_session, chunk_inputs)) + _update_retained_states(chunk_inputs, prefill_outputs) + if "image_idx_output" in prefill_outputs: + chunk_inputs["image_idx"] = prefill_outputs["image_idx_output"].astype(np.int64) + + prefill_session.deactivate() + if prefill_outputs is None: + raise RuntimeError("QAIC prefill did not execute.") + prefill_time = perf_counter() - prefill_start + vision_time + print(f"Prefill time, including vision: {prefill_time:.2f} secs") + + generated_ids = [_get_next_token_ids(prefill_outputs["logits"])] + decode_inputs = { + "input_ids": generated_ids[-1], + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, + "vision_embeds": chunk_inputs.get("vision_embeds", vision_embeds), + "image_idx": chunk_inputs.get("image_idx", np.zeros((BS, 1), dtype=np.int64)), + } + _update_retained_states(decode_inputs, prefill_outputs) + + print("Running decode QPC...") + decode_start = perf_counter() + for _ in range(1, generation_len): + decode_outputs = decode_session.run(_filter_session_inputs(decode_session, decode_inputs)) + generated_ids.append(_get_next_token_ids(decode_outputs["logits"])) + decode_inputs["input_ids"] = generated_ids[-1] + decode_inputs["position_ids"] = decode_inputs["position_ids"] + 1 + if "image_idx_output" in decode_outputs: + decode_inputs["image_idx"] = decode_outputs["image_idx_output"].astype(np.int64) + _update_retained_states(decode_inputs, decode_outputs) + + decode_time = perf_counter() - decode_start + if generation_len > 1: + print(f"Decode tok/sec: {(generation_len - 1) / decode_time:.2f}") + return np.concatenate(generated_ids, axis=1) + + +def _load_model(args): + set_deterministic(1234) + config = prepare_config(args.model_path) + kimi_cls = load_kimi_k25_class(args.model_path) + + model_kwargs = { + "config": config, + "trust_remote_code": True, + "attn_implementation": "eager", + "torch_dtype": torch.float32, + } + + if args.full_model: + model, tokenizer, processor = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) + elif args.num_vision_layers is not None and args.num_text_layers is not None: + model, tokenizer, processor = load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + ) + print( + "Loaded layer subset: " + f"vision={model.config.vision_config.vt_num_hidden_layers}, " + f"text={model.config.text_config.num_hidden_layers}, " + f"experts={model.config.text_config.n_routed_experts}" + ) + else: + raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") + + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +def _prepare_inputs(processor, args): + image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) + + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + return image, inputs + + +def main(): + args = parse_args() + model, tokenizer, processor = _load_model(args) + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + config=model.config, + torch_dtype=torch.float32, + qaic_config=qaic_config, + layerwise=False, + ) + + image, inputs = _prepare_inputs(processor, args) + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + + vision_qpc_path, prefill_qpc_path, decode_qpc_path = _compile_disagg_qpcs(qeff_model, args, image) + print(f"Vision QPC path: {vision_qpc_path}") + print(f"Prefill QPC path: {prefill_qpc_path}") + print(f"Decode QPC path: {decode_qpc_path}") + + sessions = [] + try: + vision_session = QAICInferenceSession(_resolve_qpc_path(vision_qpc_path, "vision_qpc_path")) + prefill_session = QAICInferenceSession(_resolve_qpc_path(prefill_qpc_path, "lang_prefill_qpc_path")) + decode_session = QAICInferenceSession(_resolve_qpc_path(decode_qpc_path, "lang_decode_qpc_path")) + sessions.extend([vision_session, prefill_session, decode_session]) + + generated_ids = _run_disagg_generation( + inputs, + vision_session, + prefill_session, + decode_session, + prefill_seq_len=args.prefill_seq_len, + generation_len=args.generation_len, + ) + finally: + for session in sessions: + session.deactivate() + + print(generated_ids) + print(tokenizer.batch_decode(torch.as_tensor(generated_ids), skip_special_tokens=True)) + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k2/test_kimi_k25_dynamo.py b/examples/kimi_k2/test_kimi_k25_dynamo.py new file mode 100644 index 0000000000..0ceb0da322 --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25_dynamo.py @@ -0,0 +1,429 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import importlib.util +import os +import re +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import torch +from export_kimi_k25_dynamo import ( + DEFAULT_COMPILE_DIR, + DEFAULT_EXPORT_DIR, + DEFAULT_IMAGE_URL, + compile_components, + configure_qaic_tool_path, + export_components, + get_component_export_args, + load_generation_image, + precompute_vision_rope_cache, + validate_dynamo_torch_version, +) + +from QEfficient import QEFFAutoModelForImageTextToText + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +TEXT_PROMPT = "Describe this image." +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 1024 +PREFILL_SEQ_LEN = 2 + + +def _has_qaic_runtime_access() -> bool: + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except (ImportError, OSError, RuntimeError, AttributeError): + return False + + +def _skip_test(reason: str): + try: + import pytest + + pytest.skip(reason) + except ImportError as exc: + raise RuntimeError(reason) from exc + + +def _env_bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + return default if value is None else int(value) + + +def _env_path(name: str) -> Path | None: + value = os.environ.get(name) + return None if not value else Path(value).expanduser().resolve() + + +def _parse_device_ids(value: str) -> list[int]: + return [int(device_id) for device_id in value.strip().strip("[]").split(",") if device_id.strip()] + + +def _find_free_qaic_device_id() -> int | None: + qaic_util_path = Path("/opt/qti-aic/tools/qaic-util") + if not qaic_util_path.exists(): + return None + + try: + result = subprocess.run( + [str(qaic_util_path), "-q"], + check=False, + text=True, + capture_output=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + + current_qid = None + for line in result.stdout.splitlines(): + qid_match = re.match(r"^QID\s+(\d+)", line.strip()) + if qid_match: + current_qid = int(qid_match.group(1)) + continue + + free_match = re.search(r"Nsp Free:\s*(\d+)", line) + if free_match and current_qid is not None and int(free_match.group(1)) > 0: + return current_qid + + return None + + +def _resolve_device_ids() -> list[int] | None: + value = os.environ.get("KIMI_K25_DYNAMO_DEVICE_IDS") or os.environ.get("KIMI_K25_DYNAMO_DEVICE_ID") + if value: + return _parse_device_ids(value) + + free_device_id = _find_free_qaic_device_id() + if free_device_id is None: + return None + return [free_device_id] + + +def _clone_inputs(inputs): + return {name: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for name, value in inputs.items()} + + +def _decode_tokens(tokenizer, token_ids) -> str: + decoded = tokenizer.batch_decode(torch.as_tensor(token_ids), skip_special_tokens=True) + return decoded[0] if decoded else "" + + +@torch.no_grad() +def _greedy_generate_hf(model, inputs, max_new_tokens: int) -> torch.Tensor: + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + return torch.cat(new_tokens, dim=1) + + +@torch.no_grad() +def _greedy_generate_qeff(qeff_model, inputs, args) -> torch.Tensor: + prefill_seq_len = args.prefill_seq_len + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + grid_thws = inputs["grid_thws"].to(torch.long) + + 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) + vision_embeds = qeff_model.vision_model.model( + inputs["pixel_values"].to(qeff_model.model.config.torch_dtype), + h_shape, + w_shape, + ).detach() + + input_ids_length = generated_ids.shape[1] + num_chunks = -(input_ids_length // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + generated_ids = torch.nn.functional.pad( + generated_ids, + (0, padded_len - input_ids_length), + "constant", + qeff_model.model.config.pad_token_id, + ) + attention_mask = torch.nn.functional.pad(attention_mask, (0, padded_len - input_ids_length), "constant", 0) + position_ids = torch.where( + attention_mask.bool(), + torch.arange(padded_len, dtype=torch.long).view(1, -1), + torch.full((generated_ids.shape[0], padded_len), -1, dtype=torch.long), + ) + + export_inputs, _, _ = get_component_export_args(qeff_model, prefill_seq_len) + compressed_kvs = [ + [cache_tensor.clone()[: generated_ids.shape[0]] for cache_tensor in layer_cache] + for layer_cache in export_inputs["lang"]["compressed_kvs"] + ] + image_idx = torch.zeros((generated_ids.shape[0], 1), dtype=torch.int64) + new_tokens = [] + + logits = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = start + prefill_seq_len + logits, _, image_idx_output, compressed_kvs = qeff_model.lang_model.model( + input_ids=generated_ids[:, start:end], + position_ids=position_ids[:, start:end], + vision_embeds=vision_embeds, + image_idx=image_idx, + compressed_kvs=compressed_kvs, + ) + if image_idx_output is not None: + image_idx = image_idx_output + + next_token = logits.argmax(2) + if next_token.ndim == 2 and next_token.shape[1] > 1: + next_token = next_token[:, -1:] + new_tokens.append(next_token) + + decode_position_ids = position_ids.max(dim=-1, keepdim=True).values + 1 + for _ in range(1, args.generation_len): + logits, _, image_idx_output, compressed_kvs = qeff_model.lang_model.model( + input_ids=next_token, + position_ids=decode_position_ids, + vision_embeds=vision_embeds, + image_idx=image_idx, + compressed_kvs=compressed_kvs, + ) + if image_idx_output is not None: + image_idx = image_idx_output + next_token = logits.argmax(2) + if next_token.ndim == 2 and next_token.shape[1] > 1: + next_token = next_token[:, -1:] + new_tokens.append(next_token) + decode_position_ids = decode_position_ids + 1 + + return torch.cat(new_tokens, dim=1) + + +def _build_generation_inputs(processor, image, args, dtype): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs["pixel_values"] = inputs["pixel_values"].to(dtype) + return inputs + + +def _load_hf_model(args): + model_path = args.model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + ) + return model.eval(), tokenizer, processor + + +def _build_qeff_model_from_hf(model, args): + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) + qeff_model.model.eval() + qeff_model.vision_model.model.eval() + qeff_model.lang_model.model.eval() + qeff_model.transform( + ctx_len=args.ctx_len, + seq_len=args.prefill_seq_len, + bs=1, + num_devices=args.num_devices, + qaic_config=qaic_config, + aic_num_cores=args.num_cores, + ) + precompute_vision_rope_cache(qeff_model) + return qeff_model, qaic_config + + +def _make_args(device_ids: list[int]) -> SimpleNamespace: + vision_qpc_path = _env_path("KIMI_K25_DYNAMO_VISION_QPC_PATH") + lang_qpc_path = _env_path("KIMI_K25_DYNAMO_LANG_QPC_PATH") + skip_compile = vision_qpc_path is not None and lang_qpc_path is not None + + vision_onnx_path = _env_path("KIMI_K25_DYNAMO_VISION_ONNX_PATH") + lang_onnx_path = _env_path("KIMI_K25_DYNAMO_LANG_ONNX_PATH") + skip_export = skip_compile or (vision_onnx_path is not None and lang_onnx_path is not None) + + return SimpleNamespace( + model_path=_env_path("KIMI_K25_DYNAMO_MODEL_PATH"), + export_dir=_env_path("KIMI_K25_DYNAMO_EXPORT_DIR") or DEFAULT_EXPORT_DIR, + compile_dir=_env_path("KIMI_K25_DYNAMO_COMPILE_DIR") or DEFAULT_COMPILE_DIR, + vision_onnx_path=vision_onnx_path, + lang_onnx_path=lang_onnx_path, + vision_qpc_path=vision_qpc_path, + lang_qpc_path=lang_qpc_path, + component="both", + num_vision_layers=_env_int("KIMI_K25_DYNAMO_NUM_VISION_LAYERS", NUM_VISION_LAYERS), + num_text_layers=_env_int("KIMI_K25_DYNAMO_NUM_TEXT_LAYERS", NUM_TEXT_LAYERS), + expert_ids=LOADED_EXPERT_IDS, + num_experts_per_token=NUM_EXPERTS_PER_TOKEN, + prefill_seq_len=_env_int("KIMI_K25_DYNAMO_PREFILL_SEQ_LEN", PREFILL_SEQ_LEN), + ctx_len=_env_int("KIMI_K25_DYNAMO_CTX_LEN", CTX_LEN), + num_devices=len(device_ids), + num_cores=_env_int("KIMI_K25_DYNAMO_NUM_CORES", 16), + seed=_env_int("KIMI_K25_DYNAMO_SEED", 1234), + prompt=os.environ.get("KIMI_K25_DYNAMO_PROMPT", TEXT_PROMPT), + image_url=os.environ.get("KIMI_K25_DYNAMO_IMAGE_URL", DEFAULT_IMAGE_URL), + image_path=_env_path("KIMI_K25_DYNAMO_IMAGE_PATH"), + image_height=None, + image_width=None, + generation_len=_env_int("KIMI_K25_DYNAMO_GENERATION_LEN", NEW_GENERATION_TOKENS), + device_ids=device_ids, + mxfp6_matmul=_env_bool("KIMI_K25_DYNAMO_MXFP6_MATMUL"), + mxint8_kv_cache=_env_bool("KIMI_K25_DYNAMO_MXINT8_KV_CACHE"), + mos=_env_int("KIMI_K25_DYNAMO_MOS", 1), + aic_enable_depth_first=_env_bool("KIMI_K25_DYNAMO_AIC_ENABLE_DEPTH_FIRST"), + skip_export=skip_export, + skip_compile=skip_compile, + skip_generate=False, + use_onnx_subfunctions=_env_bool("KIMI_K25_DYNAMO_USE_ONNX_SUBFUNCTIONS"), + keep_weights=_env_bool("KIMI_K25_DYNAMO_KEEP_WEIGHTS"), + ) + + +def check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + validate_dynamo_torch_version() + configure_qaic_tool_path() + if not _has_qaic_runtime_access(): + _skip_test("QAIC generation skipped: no QAIC runtime access.") + + device_ids = _resolve_device_ids() + if device_ids is None: + _skip_test("QAIC generation skipped: no QAIC device has free NSPs.") + + args = _make_args(device_ids) + set_deterministic(args.seed) + model, tokenizer, processor = _load_hf_model(args) + image = load_generation_image(args) + inputs = _build_generation_inputs(processor, image, args, model.config.torch_dtype) + + hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), args.generation_len).cpu() + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + + qeff_model, qaic_config = _build_qeff_model_from_hf(model, args) + qeff_tokens = _greedy_generate_qeff(qeff_model, _clone_inputs(inputs), args).cpu() + print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) + + assert torch.equal(hf_tokens, qeff_tokens), ( + "HF and QEff PyTorch tokens do not match for the Dynamo export wrapper path: " + f"hf={hf_tokens.tolist()}, qeff={qeff_tokens.tolist()}" + ) + + export_inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + exported_paths = export_components(qeff_model, export_inputs, output_names, dynamic_axes, args) + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + print(f"Dynamo ONNX paths: {exported_paths}") + print(f"Dynamo QPC paths: {qpc_paths}") + + qaic_output = qeff_model.generate( + inputs=_clone_inputs(inputs), + device_ids=args.device_ids, + generation_len=args.generation_len, + image_height=image.height, + image_width=image.width, + ) + qaic_tokens = torch.as_tensor(qaic_output.generated_ids[:, : args.generation_len], dtype=hf_tokens.dtype) + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens), "\n", qaic_tokens) + + if torch.equal(hf_tokens, qaic_tokens): + return + + mismatch_message = ( + "HF/QEff and QAIC tokens do not match for the Dynamo exported and compiled model: " + f"hf={hf_tokens.tolist()}, qeff={qeff_tokens.tolist()}, qaic={qaic_tokens.tolist()}" + ) + if _env_bool("KIMI_K25_DYNAMO_STRICT_HF_QAIC", False): + raise AssertionError(mismatch_message) + + print( + "QAIC token drift tolerated for this reduced Kimi K2.5 Dynamo smoke test. " + "Set KIMI_K25_DYNAMO_STRICT_HF_QAIC=1 to require exact HF-vs-QAIC token parity. " + f"{mismatch_message}" + ) + assert qaic_tokens.shape == hf_tokens.shape, "HF and QAIC generated token shapes do not match" + + +def test_kimi_k25_dynamo_hf_vs_qeff_vs_qaic(): + check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic() + + +if __name__ == "__main__": + check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic() diff --git a/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py b/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py new file mode 100644 index 0000000000..7b3d7be54b --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py @@ -0,0 +1,522 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import importlib.util +import json +import os +import re +import subprocess +from pathlib import Path +from time import perf_counter +from types import SimpleNamespace + +import onnx +import torch +from export_kimi_k25_dynamo import ( + DEFAULT_COMPILE_DIR, + DEFAULT_EXPORT_DIR, + DEFAULT_IMAGE_URL, + build_generation_inputs, + build_qeff_model, + compile_components, + configure_qaic_tool_path, + export_components, + get_component_export_args, + load_generation_image, + validate_dynamo_torch_version, +) +from onnx import numpy_helper +from safetensors.torch import save_file + +from QEfficient.generation.cloud_infer import QAICInferenceSession + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +TEXT_PROMPT = "Describe this image." +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 1024 +PREFILL_SEQ_LEN = 2 + + +class _QEffHolder: + def __init__(self, model): + self.model = model + + +def _has_qaic_runtime_access() -> bool: + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except (ImportError, OSError, RuntimeError, AttributeError): + return False + + +def _skip_test(reason: str): + try: + import pytest + + pytest.skip(reason) + except ImportError as exc: + raise RuntimeError(reason) from exc + + +def _env_bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + return default if value is None else int(value) + + +def _env_path(name: str) -> Path | None: + value = os.environ.get(name) + return None if not value else Path(value).expanduser().resolve() + + +def _env_path_any(*names: str) -> Path | None: + for name in names: + path = _env_path(name) + if path is not None: + return path + return None + + +def _env_int_any(default: int, *names: str) -> int: + for name in names: + value = os.environ.get(name) + if value is not None: + return int(value) + return default + + +def _env_bool_any(default: bool, *names: str) -> bool: + for name in names: + value = os.environ.get(name) + if value is not None: + return value.strip().lower() in {"1", "true", "yes", "on"} + return default + + +def _parse_device_ids(value: str) -> list[int]: + return [int(device_id) for device_id in value.strip().strip("[]").split(",") if device_id.strip()] + + +def _find_free_qaic_device_id() -> int | None: + candidate_paths = [Path("/opt/qti-aic/tools/qaic-util"), Path("/opt/qti-aic/exec/qaic-util")] + qaic_util_path = next((path for path in candidate_paths if path.exists()), None) + if qaic_util_path is None: + return None + + try: + result = subprocess.run( + [str(qaic_util_path), "-q"], + check=False, + text=True, + capture_output=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + + current_qid = None + for line in result.stdout.splitlines(): + qid_match = re.match(r"^QID\s+(\d+)", line.strip()) + if qid_match: + current_qid = int(qid_match.group(1)) + continue + + free_match = re.search(r"Nsp Free:\s*(\d+)", line) + if free_match and current_qid is not None and int(free_match.group(1)) > 0: + return current_qid + + return None + + +def _device_can_activate(qpc_path: Path, device_id: int) -> bool: + session = None + try: + session = QAICInferenceSession(str(qpc_path), [device_id]) + session.deactivate() + return True + except (ImportError, OSError, RuntimeError, AttributeError, ValueError): + return False + finally: + del session + + +def _find_activatable_device_id(vision_qpc_path: Path, lang_qpc_path: Path, max_device_id: int) -> int | None: + for device_id in range(max_device_id + 1): + if _device_can_activate(vision_qpc_path, device_id) and _device_can_activate(lang_qpc_path, device_id): + return device_id + return None + + +def _resolve_device_ids(args: SimpleNamespace | None = None) -> list[int] | None: + value = ( + os.environ.get("KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_IDS") + or os.environ.get("KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_ID") + or os.environ.get("KIMI_K25_DYNAMO_DEVICE_IDS") + or os.environ.get("KIMI_K25_DYNAMO_DEVICE_ID") + ) + if value: + return _parse_device_ids(value) + + if args is not None and args.vision_qpc_path is not None and args.lang_qpc_path is not None: + device_id = _find_activatable_device_id( + args.vision_qpc_path, + args.lang_qpc_path, + _env_int("KIMI_K25_DYNAMO_WEIGHT_FREE_MAX_DEVICE_ID", 63), + ) + if device_id is not None: + return [device_id] + + free_device_id = _find_free_qaic_device_id() + if free_device_id is None: + return None + return [free_device_id] + + +def _has_explicit_device_ids() -> bool: + return any( + os.environ.get(name) + for name in ( + "KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_IDS", + "KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_ID", + "KIMI_K25_DYNAMO_DEVICE_IDS", + "KIMI_K25_DYNAMO_DEVICE_ID", + ) + ) + + +def _clone_inputs(inputs): + return {name: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for name, value in inputs.items()} + + +def _decode_tokens(tokenizer, token_ids) -> str: + decoded = tokenizer.batch_decode(torch.as_tensor(token_ids), skip_special_tokens=True) + return decoded[0] if decoded else "" + + +@torch.no_grad() +def _greedy_generate_hf(model, inputs, max_new_tokens: int) -> torch.Tensor: + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + return torch.cat(new_tokens, dim=1) + + +def _load_hf_model(args): + model_path = args.model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + subset_model_path=args.export_dir / "kimi_k25_weightfree_subset" if args.use_weight_free_export else None, + ) + return model.eval(), tokenizer, processor + + +def _make_args(device_ids: list[int] | None = None) -> SimpleNamespace: + device_ids = device_ids or [0] + vision_qpc_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_VISION_QPC_PATH", "KIMI_K25_DYNAMO_VISION_QPC_PATH") + lang_qpc_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_LANG_QPC_PATH", "KIMI_K25_DYNAMO_LANG_QPC_PATH") + skip_compile = vision_qpc_path is not None and lang_qpc_path is not None + + vision_onnx_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_VISION_ONNX_PATH", "KIMI_K25_DYNAMO_VISION_ONNX_PATH") + lang_onnx_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_LANG_ONNX_PATH", "KIMI_K25_DYNAMO_LANG_ONNX_PATH") + skip_export = skip_compile or (vision_onnx_path is not None and lang_onnx_path is not None) + + return SimpleNamespace( + model_path=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_MODEL_PATH", "KIMI_K25_DYNAMO_MODEL_PATH"), + export_dir=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_EXPORT_DIR", "KIMI_K25_DYNAMO_EXPORT_DIR") + or DEFAULT_EXPORT_DIR.with_name(f"{DEFAULT_EXPORT_DIR.name}_weight_free"), + compile_dir=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_COMPILE_DIR", "KIMI_K25_DYNAMO_COMPILE_DIR") + or DEFAULT_COMPILE_DIR.with_name(f"{DEFAULT_COMPILE_DIR.name}_weight_free"), + vision_onnx_path=vision_onnx_path, + lang_onnx_path=lang_onnx_path, + vision_qpc_path=vision_qpc_path, + lang_qpc_path=lang_qpc_path, + component="both", + num_vision_layers=_env_int_any( + NUM_VISION_LAYERS, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_VISION_LAYERS", "KIMI_K25_DYNAMO_NUM_VISION_LAYERS" + ), + num_text_layers=_env_int_any( + NUM_TEXT_LAYERS, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_TEXT_LAYERS", "KIMI_K25_DYNAMO_NUM_TEXT_LAYERS" + ), + expert_ids=LOADED_EXPERT_IDS, + num_experts_per_token=NUM_EXPERTS_PER_TOKEN, + prefill_seq_len=_env_int_any( + PREFILL_SEQ_LEN, "KIMI_K25_DYNAMO_WEIGHT_FREE_PREFILL_SEQ_LEN", "KIMI_K25_DYNAMO_PREFILL_SEQ_LEN" + ), + ctx_len=_env_int_any(CTX_LEN, "KIMI_K25_DYNAMO_WEIGHT_FREE_CTX_LEN", "KIMI_K25_DYNAMO_CTX_LEN"), + num_devices=len(device_ids), + num_cores=_env_int_any(16, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_CORES", "KIMI_K25_DYNAMO_NUM_CORES"), + seed=_env_int_any(1234, "KIMI_K25_DYNAMO_WEIGHT_FREE_SEED", "KIMI_K25_DYNAMO_SEED"), + prompt=os.environ.get( + "KIMI_K25_DYNAMO_WEIGHT_FREE_PROMPT", os.environ.get("KIMI_K25_DYNAMO_PROMPT", TEXT_PROMPT) + ), + image_url=os.environ.get( + "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_URL", os.environ.get("KIMI_K25_DYNAMO_IMAGE_URL", DEFAULT_IMAGE_URL) + ), + image_path=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_PATH", "KIMI_K25_DYNAMO_IMAGE_PATH"), + image_height=_env_int_any(None, "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_HEIGHT", "KIMI_K25_DYNAMO_IMAGE_HEIGHT"), + image_width=_env_int_any(None, "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_WIDTH", "KIMI_K25_DYNAMO_IMAGE_WIDTH"), + generation_len=_env_int_any( + NEW_GENERATION_TOKENS, "KIMI_K25_DYNAMO_WEIGHT_FREE_GENERATION_LEN", "KIMI_K25_DYNAMO_GENERATION_LEN" + ), + device_ids=device_ids, + mxfp6_matmul=_env_bool_any(False, "KIMI_K25_DYNAMO_WEIGHT_FREE_MXFP6_MATMUL", "KIMI_K25_DYNAMO_MXFP6_MATMUL"), + mxint8_kv_cache=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_MXINT8_KV_CACHE", "KIMI_K25_DYNAMO_MXINT8_KV_CACHE" + ), + mos=_env_int_any(1, "KIMI_K25_DYNAMO_WEIGHT_FREE_MOS", "KIMI_K25_DYNAMO_MOS"), + aic_enable_depth_first=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_AIC_ENABLE_DEPTH_FIRST", "KIMI_K25_DYNAMO_AIC_ENABLE_DEPTH_FIRST" + ), + skip_export=skip_export, + skip_compile=skip_compile, + skip_generate=False, + use_onnx_subfunctions=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_USE_ONNX_SUBFUNCTIONS", "KIMI_K25_DYNAMO_USE_ONNX_SUBFUNCTIONS" + ), + keep_weights=True, + use_weight_free_export=True, + ) + + +def _update_onnx_weight_spec_metadata(onnx_path: Path, spec: dict): + model = onnx.load(str(onnx_path), load_external_data=False) + value = json.dumps(spec, separators=(",", ":"), sort_keys=True) + for prop in model.metadata_props: + if prop.key == "com.qti.aisw.extdata": + prop.value = value + break + else: + prop = model.metadata_props.add() + prop.key = "com.qti.aisw.extdata" + prop.value = value + tmp_path = onnx_path.with_suffix(onnx_path.suffix + ".tmp") + onnx.save(model, str(tmp_path)) + tmp_path.replace(onnx_path) + + +def _materialize_weight_free_extdata(component_model, component_dir: Path, onnx_name: str): + spec_path = component_dir / "weight_spec.json" + if not spec_path.is_file(): + raise FileNotFoundError(f"Missing weight-free spec: {spec_path}") + + spec = json.loads(spec_path.read_text()) + state_dict = component_model.state_dict() + tensors = {} + missing = [] + for entry in spec["inputs"]: + name = entry["name"] + tensor = state_dict.get(name) + if tensor is None: + missing.append(name) + continue + tensors[name] = tensor.detach().cpu().contiguous() + if missing: + raise RuntimeError( + f"Weight-free extdata tensors are missing from {component_model.__class__.__name__}: {missing}" + ) + + checkpoint_dir = component_dir / "loaded_weight_checkpoint" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / "model.safetensors" + save_file(tensors, str(checkpoint_path)) + + spec["files"] = [{"format": "safetensors", "path": "loaded_weight_checkpoint/model.safetensors"}] + for entry in spec["inputs"]: + entry["location"]["file"] = 0 + entry["location"]["key"] = entry["name"] + + spec_path.write_text(json.dumps(spec, indent=2)) + _update_onnx_weight_spec_metadata(component_dir / onnx_name, spec) + + +def _embed_packed_uint8_weights(component_dir: Path, onnx_name: str): + spec_path = component_dir / "weight_spec.json" + checkpoint_path = component_dir / "loaded_weight_checkpoint" / "model.safetensors" + spec = json.loads(spec_path.read_text()) + packed_names = { + entry["name"] + for entry in spec["inputs"] + if entry["name"].endswith("_qweight") or entry["name"].endswith("_qzeros") + } + if not packed_names: + return + + from safetensors import safe_open + + model = onnx.load(str(component_dir / onnx_name), load_external_data=False) + existing_initializers = {initializer.name for initializer in model.graph.initializer} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + for name in sorted(packed_names): + if name not in existing_initializers: + model.graph.initializer.append(numpy_helper.from_array(handle.get_tensor(name).numpy(), name=name)) + + keep_inputs = [value for value in model.graph.input if value.name not in packed_names] + del model.graph.input[:] + model.graph.input.extend(keep_inputs) + + spec["inputs"] = [entry for entry in spec["inputs"] if entry["name"] not in packed_names] + value = json.dumps(spec, separators=(",", ":"), sort_keys=True) + for prop in model.metadata_props: + if prop.key == "com.qti.aisw.extdata": + prop.value = value + break + else: + prop = model.metadata_props.add() + prop.key = "com.qti.aisw.extdata" + prop.value = value + + onnx_path = component_dir / onnx_name + tmp_path = onnx_path.with_suffix(onnx_path.suffix + ".tmp") + onnx.save(model, str(tmp_path)) + tmp_path.replace(onnx_path) + spec_path.write_text(json.dumps(spec, indent=2)) + + +def _prepare_weight_free_artifacts(qeff_model, exported_paths: dict[str, Path]): + vision_onnx_path = Path(exported_paths["vision"]) + lang_onnx_path = Path(exported_paths["lang"]) + _materialize_weight_free_extdata(qeff_model.vision_model.model, vision_onnx_path.parent, vision_onnx_path.name) + _materialize_weight_free_extdata(qeff_model.lang_model.model, lang_onnx_path.parent, lang_onnx_path.name) + _embed_packed_uint8_weights(lang_onnx_path.parent, lang_onnx_path.name) + + +def check_kimi_k25_dynamo_weight_free_hf_vs_qaic(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + validate_dynamo_torch_version() + configure_qaic_tool_path() + if not _has_qaic_runtime_access(): + _skip_test("QAIC generation skipped: no QAIC runtime access.") + + args = _make_args() + if args.skip_compile: + device_ids = _resolve_device_ids(args) + else: + device_ids = _resolve_device_ids() + if device_ids is None and not args.skip_compile: + device_ids = [0] + if device_ids is None: + _skip_test("QAIC generation skipped: no QAIC device has free NSPs.") + args = _make_args(device_ids) + + set_deterministic(args.seed) + hf_model, tokenizer, processor = _load_hf_model(args) + image = load_generation_image(args) + hf_inputs = build_generation_inputs(processor, _QEffHolder(hf_model), image, args) + hf_tokens = _greedy_generate_hf(hf_model, _clone_inputs(hf_inputs), args.generation_len).cpu() + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + + qeff_model, tokenizer, processor, qaic_config = build_qeff_model(args) + qeff_inputs = build_generation_inputs(processor, qeff_model, image, args) + export_inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + + export_start = perf_counter() + exported_paths = export_components(qeff_model, export_inputs, output_names, dynamic_axes, args) + export_end = perf_counter() + print(f"Weight-free Dynamo ONNX export time: {export_end - export_start:.2f} sec (skip_export={args.skip_export})") + + if not args.skip_compile: + prep_start = perf_counter() + _prepare_weight_free_artifacts(qeff_model, exported_paths) + prep_end = perf_counter() + print(f"Weight-free external-data prep time: {prep_end - prep_start:.2f} sec") + + compile_start = perf_counter() + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + compile_end = perf_counter() + print( + f"Weight-free Dynamo QAIC compile time: {compile_end - compile_start:.2f} sec (skip_compile={args.skip_compile})" + ) + print(f"Weight-free Dynamo ONNX paths: {exported_paths}") + print(f"Weight-free Dynamo QPC paths: {qpc_paths}") + + if not _has_explicit_device_ids(): + activatable_device_id = _find_activatable_device_id( + Path(qpc_paths["vision_qpc_path"]), + Path(qpc_paths["lang_qpc_path"]), + _env_int("KIMI_K25_DYNAMO_WEIGHT_FREE_MAX_DEVICE_ID", 63), + ) + if activatable_device_id is not None: + args.device_ids = [activatable_device_id] + print(f"Using activatable QAIC device id: {activatable_device_id}") + + qaic_output = qeff_model.generate( + inputs=_clone_inputs(qeff_inputs), + device_ids=args.device_ids, + generation_len=args.generation_len, + ) + qaic_tokens = torch.as_tensor(qaic_output.generated_ids[:, : args.generation_len], dtype=hf_tokens.dtype) + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens), "\n", qaic_tokens) + + assert torch.equal(hf_tokens, qaic_tokens), ( + "HF and QAIC tokens do not match for the weight-free Dynamo exported and compiled model: " + f"hf={hf_tokens.tolist()}, qaic={qaic_tokens.tolist()}" + ) + + +def test_kimi_k25_dynamo_weight_free_hf_vs_qaic(): + check_kimi_k25_dynamo_weight_free_hf_vs_qaic() + + +if __name__ == "__main__": + check_kimi_k25_dynamo_weight_free_hf_vs_qaic() diff --git a/tests/configs/image_text_model_configs.json b/tests/configs/image_text_model_configs.json index 63ae183c31..664aae299f 100644 --- a/tests/configs/image_text_model_configs.json +++ b/tests/configs/image_text_model_configs.json @@ -609,6 +609,57 @@ } } }, + { + "model_name": "moonshotai/Kimi-K2.5", + "model_type": "kimi_k25", + "batch_size": 1, + "prompt_len": 32, + "ctx_len": 1024, + "img_size": null, + "img_url": "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + "text_prompt": "Describe this image.", + "num_layers": 2, + "img_url_list": [ + "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png" + ], + "text_prompt_list": [ + "Describe this image.", + "Can you describe the image in detail?" + ], + "full_batch_size": 2, + "additional_params": { + "text_config": { + "hidden_size": 64, + "intermediate_size": 128, + "kv_lora_rank": 16, + "max_position_embeddings": 1024, + "model_type": "kimi_k2", + "moe_intermediate_size": 32, + "n_group": 1, + "n_routed_experts": 4, + "n_shared_experts": 1, + "num_attention_heads": 4, + "num_experts_per_tok": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 4, + "q_lora_rank": 32, + "qk_nope_head_dim": 8, + "qk_rope_head_dim": 8, + "topk_group": 1, + "v_head_dim": 16, + "vocab_size": 163840 + }, + "vision_config": { + "mm_hidden_size": 64, + "text_hidden_size": 64, + "vt_hidden_size": 64, + "vt_intermediate_size": 128, + "vt_num_attention_heads": 4, + "vt_num_hidden_layers": 2 + } + } + }, { "model_name": "tiny-random/gemma-4-dense", "model_type": "gemma4", diff --git a/tests/transformers/models/image_text_to_text/test_continuous_batching.py b/tests/transformers/models/image_text_to_text/test_continuous_batching.py index d7289dbb75..5f9ab5bff6 100644 --- a/tests/transformers/models/image_text_to_text/test_continuous_batching.py +++ b/tests/transformers/models/image_text_to_text/test_continuous_batching.py @@ -33,6 +33,13 @@ load_vlm_model_from_config, set_num_layers_vlm, ) +from tests.utils.load_kimi_utils import ( + get_kimi_k25_test_config, + is_kimi_k25, + load_kimi_k25_layer_subset_model, + load_kimi_k25_model_from_config, + run_kimi_k25_hf_model_on_pytorch_CB, +) _session = requests.Session() _session.mount("https://", HTTPAdapter(max_retries=Retry(total=3, backoff_factor=1))) @@ -59,7 +66,6 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( ): prompt_len = model_config_dict[model_name]["prompt_len"] ctx_len = model_config_dict[model_name]["ctx_len"] - max_gen_len = (NEW_GENERATION_TOKENS,) img_size = model_config_dict[model_name].get("img_size") image_urls = model_config_dict[model_name]["img_url_list"] queries = model_config_dict[model_name]["text_prompt_list"] @@ -68,7 +74,27 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( full_batch_size = model_config_dict[model_name]["full_batch_size"] max_gen_len = NEW_GENERATION_TOKENS - if config is None: + if is_kimi_k25(model_name) and config is None: + model_hf, tokenizer, processor = load_kimi_k25_layer_subset_model() + config = model_hf.config + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + torch_dtype=torch.float32, + continuous_batching=True, + ) + elif is_kimi_k25(model_name): + if config is None: + config = get_kimi_k25_test_config(model_name, model_config_dict) + model_hf, tokenizer, processor = load_kimi_k25_model_from_config(config) + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + continuous_batching=True, + ) + elif config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS ) @@ -189,6 +215,26 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, image_list, prompt_list, generation_config ) compile_kwargs["img_size"] = img_size + elif is_kimi_k25(model_name): + image_urls = [image_urls[0]] * len(queries) + for img_url in image_urls: + image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") + images.append(image) + + image_list = [images[0]] * full_batch_size + prompt_list = [queries[0]] * full_batch_size + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch_CB( + copy.deepcopy(model_hf), processor, image_list, prompt_list, max_gen_len + ) + image_height = images[0].height + image_width = images[0].width + compile_kwargs.update( + { + "prefill_seq_len": 1, + "image_height": image_height, + "image_width": image_width, + } + ) else: processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) use_fast = model_name != "mistralai/Mistral-Small-3.1-24B-Instruct-2503" @@ -231,6 +277,7 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( compile_kwargs["img_size"] = img_size qeff_model.compile(**compile_kwargs) + print("QPC Outputs (QAIC):") exec_info = qeff_model.generate( tokenizer=tokenizer, @@ -253,7 +300,12 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_CB( model_hf, images, queries, generation_config=generation_config ) else: - pytorch_hf_tokens = api_runner.run_vlm_hf_model_on_pytorch_CB(model_hf, images, queries) + if is_kimi_k25(model_name): + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch_CB( + copy.deepcopy(model_hf), processor, images, queries, max_gen_len + ) + else: + pytorch_hf_tokens = api_runner.run_vlm_hf_model_on_pytorch_CB(model_hf, images, queries) print("QPC Outputs (QAIC):") exec_info = qeff_model.generate( 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 1e932c40e9..c1e1d36510 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 @@ -36,6 +36,11 @@ load_vlm_model_from_config, set_num_layers_vlm, ) +from tests.utils.load_kimi_utils import ( + is_kimi_k25, + load_kimi_k25_layer_subset_model, + run_kimi_k25_hf_model_on_pytorch, +) from ..check_model_results import dump_and_compare_results @@ -83,7 +88,17 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( ort_tokens = None n_layer = num_hidden_layers qaic_config = copy.deepcopy(qaic_config) if qaic_config is not None else None - if config is None: + + if is_kimi_k25(model_name) and config is None: + model_hf, tokenizer, processor = load_kimi_k25_layer_subset_model() + config = model_hf.config + qeff_model = QEFFAutoModelForImageTextToText( + copy.deepcopy(model_hf), + kv_offload=kv_offload, + config=model_hf.config, + torch_dtype=torch_dtype, + ) + elif config is None: config = AutoConfig.from_pretrained( model_name, trust_remote_code=True, padding=model_name not in ModelConfig.MOLMO_MODELS ) @@ -228,6 +243,33 @@ def check_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100( inputs["pixel_values"] = inputs.pop("images") compile_kwargs["img_size"] = img_size + elif is_kimi_k25(model_name): + image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + prompt = processor.apply_chat_template(conversation, add_generation_prompt=True) + inputs = processor( + messages=conversation, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + pytorch_hf_tokens = run_kimi_k25_hf_model_on_pytorch(copy.deepcopy(model_hf), processor, inputs, max_gen_len) + compile_kwargs.update( + { + "prefill_seq_len": 1, + "image_height": image.height, + "image_width": image.width, + } + ) + else: processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, padding=True) image = Image.open(_session.get(img_url, stream=True).raw) @@ -486,6 +528,7 @@ def test_image_text_to_text_pytorch_vs_kv_vs_ort_vs_ai100_qnn(model_name, kv_off "google/gemma-3-4b-it", "tiny-random/gemma-4-dense", "tiny-random/gemma-4-moe", + "moonshotai/Kimi-K2.5", ]: pytest.skip("QNN is not supported for these models yet.") diff --git a/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py new file mode 100644 index 0000000000..05507e8f1f --- /dev/null +++ b/tests/transformers/models/image_text_to_text/test_kimi_k25_disagg.py @@ -0,0 +1,319 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +import copy +from io import BytesIO +from pathlib import Path + +import numpy as np +import pytest +import requests +import torch +from PIL import Image + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession +from tests.utils.load_kimi_utils import ( + LOADED_EXPERT_IDS, + NUM_EXPERTS_PER_TOKEN, + load_kimi_k25_class, + load_layer_subset_model, + prepare_config, + resolve_model_path, + run_kimi_k25_hf_model_on_pytorch, + set_deterministic, +) + +PREFILL_SEQ_LEN = 64 +CTX_LEN = 512 +BATCH_SIZE = 1 +GENERATION_LEN = 4 +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" +IMAGE_SIZE = 448 +TEXT_PROMPT = "Describe this image." + + +def _prepare_inputs(processor): + image = Image.open(BytesIO(requests.get(IMAGE_URL, timeout=30).content)).convert("RGB") + image = image.resize((IMAGE_SIZE, IMAGE_SIZE)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": TEXT_PROMPT}, + ], + } + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + return inputs, image.height, image.width + + +def _decode_tokens(tokenizer, token_ids: torch.Tensor) -> str: + decoded = tokenizer.batch_decode(token_ids, skip_special_tokens=True) + return decoded[0] if decoded else "" + + +def _clone_inputs(inputs): + return {k: (v.clone() if torch.is_tensor(v) else copy.deepcopy(v)) for k, v in inputs.items()} + + +def _assert_onnx_path(onnx_path, label: str) -> Path: + assert onnx_path is not None, f"{label} compile did not set an ONNX path" + onnx_path = Path(onnx_path) + assert onnx_path.is_file(), f"{label} ONNX path does not exist: {onnx_path}" + assert onnx_path.suffix == ".onnx", f"{label} path is not an ONNX file: {onnx_path}" + return onnx_path.resolve() + + +def _assert_distinct_onnx_paths(onnx_paths: dict[str, Path]): + unique_paths = {str(path) for path in onnx_paths.values()} + assert len(unique_paths) == len(onnx_paths), f"Expected distinct ONNX paths per compile, got: {onnx_paths}" + + +def _numpy(value): + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _session_input_names(session: QAICInferenceSession) -> set[str]: + input_names = set(session.input_names) + input_names.update(name.rsplit("/", 1)[-1] for name in session.input_names) + return input_names + + +def _cast_for_session(session: QAICInferenceSession, name: str, value: np.ndarray) -> np.ndarray: + binding_index = session.binding_index_map.get(name) + if binding_index is None: + return value + dtype = session.aic_to_np_dtype_mapping[session.bindings[binding_index].type] + return value.astype(dtype, copy=False) + + +def _filter_session_inputs(session: QAICInferenceSession, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + input_names = _session_input_names(session) + return {name: _cast_for_session(session, name, value) for name, value in inputs.items() if name in input_names} + + +def _get_next_token_ids(logits: np.ndarray) -> np.ndarray: + logits = np.asarray(logits) + return logits[:, -1, :].argmax(axis=-1).astype(np.int64).reshape(BATCH_SIZE, 1) + + +def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs: dict[str, np.ndarray]): + for output_name, value in source_outputs.items(): + output_basename = output_name.rsplit("/", 1)[-1] + if not output_basename.endswith("_RetainedState"): + continue + target_inputs[output_basename.removesuffix("_RetainedState")] = value + + +def _load_kimi_subset_model(): + set_deterministic(1234) + model_path = resolve_model_path() + config = prepare_config(model_path) + kimi_cls = load_kimi_k25_class(model_path) + + model, tokenizer, processor = load_layer_subset_model( + model_path=model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=NUM_VISION_LAYERS, + num_text_layers=NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok=NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + ) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +def _get_image_compile_dims(image_height: int, image_width: int) -> dict[str, int]: + return {"image_height": image_height, "image_width": image_width} + + +def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, compile_dims: dict[str, int]): + common_compile_kwargs = { + "batch_size": BATCH_SIZE, + "ctx_len": CTX_LEN, + "num_cores": 16, + "mxfp6_matmul": False, + "split_model_io": True, + "mos": 1, + "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, + "layerwise": False, + **compile_dims, + } + + compiled_onnx_paths = {} + vision_qpc_path = qeff_model.compile( + prefill_seq_len=PREFILL_SEQ_LEN, + skip_vision=False, + skip_lang=True, + num_devices=1, + **common_compile_kwargs, + ) + compiled_onnx_paths["vision"] = _assert_onnx_path(qeff_model.vision_model.onnx_path, "vision") + + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=PREFILL_SEQ_LEN, + prefill_only=True, + skip_vision=True, + skip_lang=False, + num_devices=1, + **common_compile_kwargs, + ) + compiled_onnx_paths["prefill"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "prefill") + + decode_qpc_path = qeff_model.compile( + prefill_seq_len=1, + prefill_only=False, + skip_vision=True, + skip_lang=False, + num_devices=1, + **common_compile_kwargs, + ) + compiled_onnx_paths["decode"] = _assert_onnx_path(qeff_model.lang_model.onnx_path, "decode") + _assert_distinct_onnx_paths(compiled_onnx_paths) + + return vision_qpc_path, prefill_qpc_path, decode_qpc_path, compiled_onnx_paths + + +def _run_disagg_qaic_generation( + common_inputs: dict[str, torch.Tensor], + vision_session: QAICInferenceSession, + prefill_session: QAICInferenceSession, + decode_session: QAICInferenceSession, +) -> np.ndarray: + inputs = {name: _numpy(value) for name, value in _clone_inputs(common_inputs).items()} + input_ids_length = inputs["input_ids"].shape[1] + num_chunks = -(input_ids_length // -PREFILL_SEQ_LEN) + padded_len = num_chunks * PREFILL_SEQ_LEN + + inputs["input_ids"] = np.pad( + inputs["input_ids"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=1, + ) + inputs["attention_mask"] = np.pad( + inputs["attention_mask"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=0, + ) + + grid_thws = inputs["grid_thws"].astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + } + vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) + vision_session.deactivate() + + vision_embeds = vision_outputs.get("vision_embeds") + assert vision_embeds is not None, f"Vision QPC did not return vision_embeds. Outputs: {vision_outputs.keys()}" + + lang_inputs = { + "input_ids": inputs["input_ids"].astype(np.int64), + "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), + "vision_embeds": vision_embeds, + "image_idx": np.zeros((BATCH_SIZE, 1), dtype=np.int64), + } + + prefill_session.set_buffers(vision_outputs) + chunk_inputs = lang_inputs.copy() + prefill_outputs = None + for chunk_idx in range(num_chunks): + start = chunk_idx * PREFILL_SEQ_LEN + end = (chunk_idx + 1) * PREFILL_SEQ_LEN + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, start:end] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][:, start:end] + prefill_outputs = prefill_session.run(_filter_session_inputs(prefill_session, chunk_inputs)) + _update_retained_states(chunk_inputs, prefill_outputs) + if "image_idx_output" in prefill_outputs: + chunk_inputs["image_idx"] = prefill_outputs["image_idx_output"].astype(np.int64) + + prefill_session.deactivate() + assert prefill_outputs is not None, "QAIC prefill did not execute." + + generated_ids = [_get_next_token_ids(prefill_outputs["logits"])] + decode_inputs = { + "input_ids": generated_ids[-1], + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, + "vision_embeds": chunk_inputs.get("vision_embeds", vision_embeds), + "image_idx": chunk_inputs.get("image_idx", np.zeros((BATCH_SIZE, 1), dtype=np.int64)), + } + _update_retained_states(decode_inputs, prefill_outputs) + + for _ in range(1, GENERATION_LEN): + decode_outputs = decode_session.run(_filter_session_inputs(decode_session, decode_inputs)) + generated_ids.append(_get_next_token_ids(decode_outputs["logits"])) + decode_inputs["input_ids"] = generated_ids[-1] + decode_inputs["position_ids"] = decode_inputs["position_ids"] + 1 + if "image_idx_output" in decode_outputs: + decode_inputs["image_idx"] = decode_outputs["image_idx_output"].astype(np.int64) + _update_retained_states(decode_inputs, decode_outputs) + + return np.concatenate(generated_ids, axis=1)[0] + + +@pytest.mark.on_qaic +@pytest.mark.multimodal +def test_kimi_k25_disagg_qaic_vs_hf_fp32(): + model, tokenizer, processor = _load_kimi_subset_model() + inputs, image_height, image_width = _prepare_inputs(processor) + inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + hf_tokens = run_kimi_k25_hf_model_on_pytorch( + copy.deepcopy(model), processor, _clone_inputs(inputs), max_gen_len=GENERATION_LEN + ) + + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + config=model.config, + torch_dtype=torch.float32, + layerwise=False, + ) + + compile_dims = _get_image_compile_dims(image_height, image_width) + vision_qpc_path, prefill_qpc_path, decode_qpc_path, compiled_onnx_paths = _compile_disagg_qpcs( + qeff_model, + compile_dims, + ) + print(f"Kimi-K2.5 disagg ONNX paths: {compiled_onnx_paths}") + + sessions = [] + try: + vision_session = QAICInferenceSession(vision_qpc_path.get("vision_qpc_path")) + prefill_session = QAICInferenceSession(prefill_qpc_path.get("lang_prefill_qpc_path")) + decode_session = QAICInferenceSession(decode_qpc_path.get("lang_decode_qpc_path")) + sessions.extend([vision_session, prefill_session, decode_session]) + qaic_tokens = _run_disagg_qaic_generation( + common_inputs=inputs, + vision_session=vision_session, + prefill_session=prefill_session, + decode_session=decode_session, + ) + finally: + for session in sessions: + session.deactivate() + + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + print("Disagg QAIC:", _decode_tokens(tokenizer, torch.as_tensor(qaic_tokens)), "\n", qaic_tokens) + + assert torch.equal(hf_tokens, torch.as_tensor(qaic_tokens)), "HF and disagg QAIC tokens do not match" 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 9cb6d5f273..aca6aa3245 100644 --- a/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py +++ b/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py @@ -171,7 +171,10 @@ 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) + export_kwargs = {"use_onnx_subfunctions": subfunc, "prefill_only": prefill_only} + if prefill_only: + export_kwargs["prefill_seq_len"] = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + qeff_model.export(tmp_path, **export_kwargs) hash_params = {} hash_params["config"] = qeff_model.model.config.to_diff_dict() hash_params["peft_config"] = None diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 72728b65eb..5a362c09f3 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -18,6 +18,7 @@ but do not yet have a stable CPU runtime parity path in the consolidated test """ +import json import logging import os import shutil @@ -33,6 +34,7 @@ import onnxruntime as ort import pytest import torch +from torch import nn from transformers import ( AutoConfig, AutoModel, @@ -927,6 +929,125 @@ def _assert_qwen_hf_qeff_ort_parity(model_type: str, tmp_path, *, prefill_only: assert np.allclose(qeff_logits, ort_logits, atol=atol, rtol=1e-4) +def _kimi_k25_prompt_inputs(processor): + from PIL import Image + + image = Image.new("RGB", (64, 64), color=(128, 64, 32)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": "Describe."}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + return {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + + +def _kimi_k25_qeff_lang_inputs(qeff_model, inputs, vision_embeds): + seq_len = inputs["input_ids"].shape[1] + lang_inputs = deepcopy(qeff_model.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=seq_len)["lang"]) + lang_inputs["input_ids"] = inputs["input_ids"].clone() + lang_inputs["position_ids"] = torch.where( + inputs["attention_mask"] > 0, + torch.arange(seq_len, dtype=torch.int64).view(1, seq_len), + -1, + ) + lang_inputs["vision_embeds"] = vision_embeds + lang_inputs["image_idx"] = torch.zeros((1, 1), dtype=torch.int64) + return lang_inputs + + +class _KimiSyntheticQuantLinear(nn.Module): + def __init__(self, linear: nn.Linear): + super().__init__() + if linear.in_features % 2 != 0: + raise ValueError("Synthetic int4 fixture requires an even input dimension") + + self.in_features = linear.in_features + self.out_features = linear.out_features + self.bits = 4 + self.group_size = 1 + self.act_order = None + self.register_buffer( + "qweight", torch.full((linear.out_features, linear.in_features // 2), 0x99, dtype=torch.uint8) + ) + self.register_buffer( + "qzeros", torch.full((linear.out_features, linear.in_features // 2), 0x88, dtype=torch.uint8) + ) + self.register_buffer("scales", linear.weight.detach().clone().to(torch.float32)) + self.register_buffer("g_idx", torch.arange(linear.in_features, dtype=torch.int32)) + if linear.bias is None: + self.bias = None + else: + self.register_buffer("bias", linear.bias.detach().clone().to(torch.float32)) + + def forward(self, inputs): + output = torch.matmul(inputs.float(), self.scales.transpose(0, 1).float()) + if self.bias is not None: + output = output + self.bias.to(output.dtype) + return output.to(inputs.dtype) + + +def _install_kimi_synthetic_quant_experts(model): + for module in model.modules(): + experts = getattr(module, "experts", None) + if experts is None: + continue + for expert in experts: + for projection_name in ("gate_proj", "up_proj", "down_proj"): + projection = getattr(expert, projection_name, None) + if isinstance(projection, nn.Linear): + setattr(expert, projection_name, _KimiSyntheticQuantLinear(projection)) + + +@pytest.mark.llm_model +def test_kimi_k25_quickcheck_hf_qeff_vision_logits_parity(): + from tests.utils.load_kimi_utils import get_kimi_k25_test_config, load_kimi_k25_model_from_config + + model_id = "moonshotai/Kimi-K2.5" + config_path = Path(__file__).parents[2] / "configs" / "image_text_model_configs.json" + model_configs = json.loads(config_path.read_text())["image_text_models"] + model_config_dict = {model["model_name"]: model for model in model_configs} + + try: + config = get_kimi_k25_test_config(model_id, model_config_dict) + model_hf, _, processor = load_kimi_k25_model_from_config(config) + _install_kimi_synthetic_quant_experts(model_hf) + except Exception as exc: + _skip_on_model_fetch_error(exc, model_id) + + inputs = _kimi_k25_prompt_inputs(processor) + with torch.no_grad(): + hf_outputs = model_hf(**inputs, use_cache=False, return_dict=True) + hf_logits = hf_outputs.logits[:, -1:, :].detach().float().numpy() + + qeff_model = QEFFAutoModelForImageTextToText( + deepcopy(model_hf), + kv_offload=True, + config=model_hf.config, + torch_dtype=torch.float32, + ) + 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) + + with torch.no_grad(): + vision_embeds = qeff_model.vision_model.model(inputs["pixel_values"], h_shape, w_shape) + lang_inputs = _kimi_k25_qeff_lang_inputs(qeff_model, inputs, vision_embeds) + qeff_logits = qeff_model.lang_model.model(**lang_inputs)[0].detach().float().numpy() + + assert qeff_logits.shape == hf_logits.shape + assert np.allclose(hf_logits, qeff_logits, atol=1e-4, rtol=1e-4) + + @pytest.mark.llm_model @pytest.mark.parametrize( ("model_type", "model_id"), @@ -3551,3 +3672,50 @@ def test_layerwise_export_default_names_unchanged(tmp_path): assert f"past_key.{window}" in captured["input_names"] assert all("_vllmKvCache" not in n and "_VLLM" not in n for n in captured["output_names"]) assert all("_vllmKvCache" not in n and "_VLLM" not in n for n in captured["input_names"]) + + +def test_kimi_k25_get_specializations_supports_multi_resolution_grid_sizes(): + """Kimi K2.5 accepts list-valued image sizes for multi-resolution specs.""" + from types import SimpleNamespace + + from QEfficient.transformers.models.kimi_k25.modeling_kimi_k25 import QEffKimiK25ForConditionalGeneration + + model = QEffKimiK25ForConditionalGeneration.__new__(QEffKimiK25ForConditionalGeneration) + model.config = SimpleNamespace(vision_config=SimpleNamespace(patch_size=14, merge_kernel_size=(2, 2))) + + specs, _ = model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + image_height=[512, 448], + image_width=[910, 448], + num_frames=[1, 1], + kv_offload=True, + ) + assert specs["vision"] == [ + {"num_patches": 2508, "grid_h": 38, "grid_w": 66, "num_image_tokens": 627}, + {"num_patches": 1024, "grid_h": 32, "grid_w": 32, "num_image_tokens": 256}, + ] + assert all(spec["num_image_tokens"] == 627 for spec in specs["lang"]) + + with pytest.raises(ValueError, match="image_height and image_width"): + model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + h=[30, 32], + w=[80, 64], + num_frames=[1, 2], + kv_offload=True, + ) + + with pytest.raises(ValueError, match="num_patches"): + model.get_specializations( + batch_size=1, + prefill_seq_len=64, + ctx_len=4096, + image_height=512, + image_width=910, + num_patches=2508, + kv_offload=True, + ) diff --git a/tests/unit_test/transforms/test_quantization_transforms.py b/tests/unit_test/transforms/test_quantization_transforms.py index b7fa03c1d2..e89432873c 100644 --- a/tests/unit_test/transforms/test_quantization_transforms.py +++ b/tests/unit_test/transforms/test_quantization_transforms.py @@ -93,6 +93,7 @@ def test_all_transforms_have_mutate_classmethod(self): FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) for cls in [ @@ -100,6 +101,7 @@ def test_all_transforms_have_mutate_classmethod(self): GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ]: assert hasattr(cls, "mutate"), f"{cls.__name__} missing mutate method" assert callable(cls.mutate), f"{cls.__name__}.mutate is not callable" @@ -111,6 +113,7 @@ def test_all_transforms_are_subclasses_of_module_mutator(self): FP8DeQuantLinearToLinearTransform, GPTQToMatmulNbitsTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) for cls in [ @@ -118,6 +121,7 @@ def test_all_transforms_are_subclasses_of_module_mutator(self): GPTQToMatmulNbitsTransform, FP8DeQuantLinearToLinearTransform, Mxfp4GptOssExpertDequantizeTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ]: assert issubclass(cls, ModuleMutatorTransform), ( f"{cls.__name__} must be a subclass of ModuleMutatorTransform" @@ -328,6 +332,57 @@ def test_quantization_transforms_come_before_kv_cache_transform(self): f"AwqToMatmulNbitsTransform (idx={awq_idx}) must come before KVCacheTransform (idx={kv_idx})" ) + def test_image_text_wrappers_include_pack_quantized_int4_transform(self): + from QEfficient.transformers.models.modeling_auto import ( + QEffCausalLMForTextImageToTextModel, + QEffVisionEncoderForTextImageToTextModel, + ) + from QEfficient.transformers.models.pytorch_transforms import CustomOpsTransform + from QEfficient.transformers.quantizers.quant_transforms import PackQuantizedInt4ToMatMulNBitsTransform + + for wrapper_cls in [QEffVisionEncoderForTextImageToTextModel, QEffCausalLMForTextImageToTextModel]: + transforms = wrapper_cls._pytorch_transforms + pack_idx = transforms.index(PackQuantizedInt4ToMatMulNBitsTransform) + custom_ops_idx = transforms.index(CustomOpsTransform) + assert pack_idx < custom_ops_idx + + def test_pack_quantized_int4_transform_matches_compressed_linear_metadata(self, monkeypatch): + import importlib + from types import SimpleNamespace + + import torch + + from QEfficient.customop.matmulnbits import QuantLinearORT + from QEfficient.transformers.quantizers.quant_transforms import PackQuantizedInt4ToMatMulNBitsTransform + + linear = torch.nn.Linear(8, 8, bias=False) + del linear._parameters["weight"] + linear.register_parameter( + "weight_packed", torch.nn.Parameter(torch.zeros((8, 4), dtype=torch.int32), requires_grad=False) + ) + linear.register_parameter( + "weight_scale", torch.nn.Parameter(torch.ones((8, 2), dtype=torch.float32), requires_grad=False) + ) + linear.quantization_scheme = SimpleNamespace( + weights=SimpleNamespace(num_bits=4, type="int", strategy="group", group_size=4) + ) + + def fake_decompress_module(module): + module.weight = torch.nn.Parameter( + torch.zeros(module.out_features, module.in_features), requires_grad=False + ) + + compressor_base = importlib.import_module("compressed_tensors.compressors.base") + monkeypatch.setattr(compressor_base, "decompress_module", fake_decompress_module) + + model = torch.nn.Sequential(linear) + transformed_model, transformed = PackQuantizedInt4ToMatMulNBitsTransform.apply(model) + + assert transformed + assert isinstance(transformed_model[0], QuantLinearORT) + assert transformed_model[0].bits == 4 + assert transformed_model[0].group_size == 4 + def test_non_quantized_model_not_affected_by_quant_transforms(self): """Applying quantization transforms to a non-quantized model must not change it.""" import torch @@ -336,6 +391,7 @@ def test_non_quantized_model_not_affected_by_quant_transforms(self): from QEfficient.transformers.quantizers.quant_transforms import ( AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, + PackQuantizedInt4ToMatMulNBitsTransform, ) cfg = GPT2Config(n_layer=1, n_head=2, n_embd=64, vocab_size=500, n_positions=32, n_ctx=32) @@ -349,9 +405,16 @@ def test_non_quantized_model_not_affected_by_quant_transforms(self): model_gptq, applied_gptq = GPTQToMatmulNbitsTransform.apply(model) assert not applied_gptq, "GPTQToMatmulNbitsTransform must not apply to non-quantized model" + model_pack, applied_pack = PackQuantizedInt4ToMatMulNBitsTransform.apply(model) + assert not applied_pack, "PackQuantizedInt4ToMatMulNBitsTransform must not apply to non-quantized model" + # Model output must be unchanged input_ids = torch.randint(0, 500, (1, 8)) with torch.no_grad(): original_logits = model(input_ids=input_ids).logits awq_logits = model_awq(input_ids=input_ids).logits + pack_logits = model_pack(input_ids=input_ids).logits assert torch.allclose(original_logits, awq_logits), "AWQ transform must not change non-quantized model output" + assert torch.allclose(original_logits, pack_logits), ( + "Packed-int4 transform must not change non-quantized model output" + ) diff --git a/tests/utils/load_kimi_utils.py b/tests/utils/load_kimi_utils.py new file mode 100644 index 0000000000..f7c1431757 --- /dev/null +++ b/tests/utils/load_kimi_utils.py @@ -0,0 +1,512 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import inspect +import json +import os +import random +import re +import shutil +import sys +import tempfile +from pathlib import Path + +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoProcessor, AutoTokenizer +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +KIMI_K25_MODEL_NAME = "moonshotai/Kimi-K2.5" +NUM_VISION_LAYERS = 2 +NUM_TEXT_LAYERS = 2 +LOADED_EXPERT_IDS = (0, 1, 2, 3) +NUM_EXPERTS_PER_TOKEN = 2 + +EXPERT_KEY_PATTERN = re.compile(r"^(language_model\.model\.layers\.\d+\.mlp\.experts\.)(\d+)(\..+)$") + + +def is_kimi_k25(model_name: str) -> bool: + return model_name == KIMI_K25_MODEL_NAME + + +def set_deterministic(seed: int): + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + + +def resolve_model_path(model_name: str = KIMI_K25_MODEL_NAME) -> Path: + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + return Path(snapshot_download(repo_id=model_name, cache_dir=os.environ.get("HF_HUB_CACHE"))) + + +def patch_kimi_tie_weights_compat(kimi_cls): + tie_signature = inspect.signature(kimi_cls.tie_weights) + if tuple(tie_signature.parameters) != ("self",): + return + + def _tie_weights_compat(self, missing_keys=None, recompute_mapping=True): + lm_tie_weights = self.language_model.tie_weights + try: + return lm_tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping) + except TypeError: + return lm_tie_weights() + + kimi_cls.tie_weights = _tie_weights_compat + + +def patch_deepseek_init_weights_compat(kimi_cls): + module_prefix, _ = kimi_cls.__module__.rsplit(".", maxsplit=1) + deepseek_module = sys.modules.get(f"{module_prefix}.modeling_deepseek") + if deepseek_module is None or not hasattr(deepseek_module, "DeepseekV3PreTrainedModel"): + return + + deepseek_cls = deepseek_module.DeepseekV3PreTrainedModel + if ( + getattr(deepseek_cls, "_qeff_kimi_k25_init_weights_patched", False) + or getattr(deepseek_cls, "_qeff_test_init_weights_patched", False) + or getattr(deepseek_cls, "_qeff_t55_init_weights_patched", False) + ): + return + + def _init_weights_compat(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "bias") and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + if hasattr(module, "weight") and module.weight is not None: + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + deepseek_cls._init_weights = _init_weights_compat + deepseek_cls._qeff_kimi_k25_init_weights_patched = True + deepseek_cls._qeff_test_init_weights_patched = True + deepseek_cls._qeff_t55_init_weights_patched = True + + +def load_kimi_k25_class(model_path_or_name): + kimi_cls = get_class_from_dynamic_module( + "modeling_kimi_k25.KimiK25ForConditionalGeneration", + str(model_path_or_name), + ) + patch_kimi_tie_weights_compat(kimi_cls) + patch_deepseek_init_weights_compat(kimi_cls) + return kimi_cls + + +def patch_kimi_k25_remote_code_compat(config): + return load_kimi_k25_class(config._name_or_path) + + +def prepare_config(model_path: Path): + config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + + config._attn_implementation = "eager" + if hasattr(config, "text_config"): + config.text_config._attn_implementation = "eager" + if hasattr(config, "vision_config"): + config.vision_config._attn_implementation = "eager" + return config + + +def get_kimi_k25_test_config(model_name: str, model_config_dict): + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + config._attn_implementation = "eager" + config.torch_dtype = torch.float32 + config.dtype = torch.float32 + additional_params = model_config_dict[model_name]["additional_params"] + + for attr, value in additional_params["text_config"].items(): + setattr(config.text_config, attr, value) + config.text_config._attn_implementation = "eager" + config.text_config.torch_dtype = torch.float32 + config.text_config.dtype = torch.float32 + + for attr, value in additional_params["vision_config"].items(): + setattr(config.vision_config, attr, value) + config.vision_config._attn_implementation = "eager" + config.vision_config.torch_dtype = torch.float32 + config.vision_config.dtype = torch.float32 + + patch_kimi_k25_remote_code_compat(config) + return config + + +def load_kimi_k25_model_from_config(config): + kimi_cls = patch_kimi_k25_remote_code_compat(config) + model = kimi_cls._from_config(config) + torch_dtype = getattr(model.config, "torch_dtype", None) + if torch_dtype == torch.bfloat16 or torch_dtype == torch.float16: + model = model.to(torch.float32) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + model.eval() + tokenizer = AutoTokenizer.from_pretrained(KIMI_K25_MODEL_NAME, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(KIMI_K25_MODEL_NAME, trust_remote_code=True) + return model, tokenizer, processor + + +def get_kimi_k25_num_image_tokens(config, grid_thws): + merge_height, merge_width = config.vision_config.merge_kernel_size + return int(grid_thws[0, 1].item() // merge_height) * int(grid_thws[0, 2].item() // merge_width) + + +def parse_expert_ids(value: str): + expert_ids = tuple(int(expert_id) for expert_id in value.split(",") if expert_id.strip()) + if not expert_ids: + raise argparse.ArgumentTypeError("At least one expert id must be provided.") + return expert_ids + + +def _validate_layer_count(name, requested_count, available_count): + if requested_count < 1: + raise ValueError(f"{name} must be >= 1, got {requested_count}.") + if requested_count > available_count: + raise ValueError(f"{name}={requested_count} exceeds available layers={available_count}.") + + +def _validate_expert_subset(loaded_expert_ids, num_experts_per_tok, total_experts): + expert_ids = tuple(loaded_expert_ids) + if len(expert_ids) != 4: + raise ValueError(f"Expected exactly 4 routed experts, got {expert_ids!r}.") + if len(set(expert_ids)) != len(expert_ids): + raise ValueError(f"Expert ids must be unique, got {expert_ids!r}.") + invalid_ids = [expert_id for expert_id in expert_ids if expert_id < 0 or expert_id >= total_experts] + if invalid_ids: + raise ValueError(f"Expert ids {invalid_ids!r} are outside the valid range [0, {total_experts - 1}].") + if num_experts_per_tok > len(expert_ids): + raise ValueError(f"num_experts_per_tok={num_experts_per_tok} cannot exceed {len(expert_ids)} loaded experts.") + return expert_ids + + +def _remap_checkpoint_key(checkpoint_key, expert_index_map): + match = EXPERT_KEY_PATTERN.match(checkpoint_key) + if not match: + return checkpoint_key + + original_expert_idx = int(match.group(2)) + remapped_expert_idx = expert_index_map.get(original_expert_idx) + if remapped_expert_idx is None: + return None + return f"{match.group(1)}{remapped_expert_idx}{match.group(3)}" + + +def _is_routed_gate_weight(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.weight") + + +def _is_routed_gate_bias(checkpoint_key): + return checkpoint_key.endswith(".mlp.gate.e_score_correction_bias") + + +def allowed_prefixes(num_vision_layers: int, num_text_layers: int): + prefixes = [ + "vision_tower.patch_embed.", + "vision_tower.encoder.final_layernorm.", + "mm_projector.", + "language_model.model.embed_tokens.", + "language_model.model.norm.", + "language_model.lm_head.", + ] + prefixes.extend(f"vision_tower.encoder.blocks.{layer_idx}." for layer_idx in range(num_vision_layers)) + prefixes.extend(f"language_model.model.layers.{layer_idx}." for layer_idx in range(num_text_layers)) + return prefixes + + +def build_layer_subset_config(config, num_vision_layers, num_text_layers, loaded_expert_ids, num_experts_per_tok): + stripped_config = copy.deepcopy(config) + text_config = stripped_config.text_config + vision_config = stripped_config.vision_config + + _validate_layer_count("num_text_layers", num_text_layers, text_config.num_hidden_layers) + _validate_layer_count("num_vision_layers", num_vision_layers, vision_config.vt_num_hidden_layers) + + text_config.num_hidden_layers = num_text_layers + vision_config.vt_num_hidden_layers = num_vision_layers + + loaded_expert_ids = _validate_expert_subset( + loaded_expert_ids, + num_experts_per_tok, + text_config.n_routed_experts, + ) + text_config.n_routed_experts = len(loaded_expert_ids) + text_config.num_experts_per_tok = num_experts_per_tok + text_config.n_group = 1 + text_config.topk_group = 1 + return stripped_config, loaded_expert_ids + + +def materialize_subset_checkpoint( + model_path: Path, temp_model_path: Path, weight_map, allowed_weight_prefixes, loaded_expert_ids +): + expert_index_map = {expert_id: remapped_idx for remapped_idx, expert_id in enumerate(loaded_expert_ids)} + shard_to_entries = {} + for checkpoint_key, source_shard_name in weight_map.items(): + if not checkpoint_key.startswith(tuple(allowed_weight_prefixes)): + continue + + remapped_key = _remap_checkpoint_key(checkpoint_key, expert_index_map) + if remapped_key is None: + continue + shard_to_entries.setdefault(source_shard_name, []).append((checkpoint_key, remapped_key)) + + filtered_weight_map = {} + subset_shards = [] + for shard_idx, (source_shard_name, shard_entries) in enumerate(sorted(shard_to_entries.items())): + tensors = {} + with safe_open(model_path / source_shard_name, framework="pt", device="cpu") as shard_reader: + for checkpoint_key, remapped_key in shard_entries: + tensor = shard_reader.get_tensor(checkpoint_key) + if _is_routed_gate_weight(checkpoint_key): + tensor = tensor[list(loaded_expert_ids), :].contiguous() + elif _is_routed_gate_bias(checkpoint_key): + tensor = tensor[list(loaded_expert_ids)].contiguous() + tensors[remapped_key] = tensor + + subset_shard_name = f"model-subset-{shard_idx:05d}.safetensors" + save_file(tensors, str(temp_model_path / subset_shard_name)) + subset_shards.append(subset_shard_name) + filtered_weight_map.update({remapped_key: subset_shard_name for _, remapped_key in shard_entries}) + + return filtered_weight_map, subset_shards + + +def load_layer_subset_model( + *, + model_path: Path, + kimi_cls, + config, + num_vision_layers: int, + num_text_layers: int, + loaded_expert_ids, + num_experts_per_tok: int, + dtype, + subset_model_path: Path | None = None, +): + checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) + weight_map = checkpoint_index["weight_map"] + stripped_config, loaded_expert_ids = build_layer_subset_config( + config, + num_vision_layers=num_vision_layers, + num_text_layers=num_text_layers, + loaded_expert_ids=loaded_expert_ids, + num_experts_per_tok=num_experts_per_tok, + ) + + tempdir_context = None + if subset_model_path is None: + tempdir_context = tempfile.TemporaryDirectory() + temp_model_path = Path(tempdir_context.__enter__()) + else: + temp_model_path = Path(subset_model_path).expanduser().resolve() + if temp_model_path.exists(): + shutil.rmtree(temp_model_path) + temp_model_path.mkdir(parents=True, exist_ok=True) + + try: + filtered_weight_map, subset_shards = materialize_subset_checkpoint( + model_path=model_path, + temp_model_path=temp_model_path, + weight_map=weight_map, + allowed_weight_prefixes=allowed_prefixes(num_vision_layers, num_text_layers), + loaded_expert_ids=loaded_expert_ids, + ) + (temp_model_path / "config.json").write_text(stripped_config.to_json_string(use_diff=False)) + (temp_model_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": { + "total_size": sum((temp_model_path / shard_name).stat().st_size for shard_name in subset_shards) + }, + "weight_map": filtered_weight_map, + } + ) + ) + + model_kwargs = { + "config": stripped_config, + "trust_remote_code": True, + "attn_implementation": "eager", + "output_loading_info": True, + } + if dtype is not None: + model_kwargs["torch_dtype"] = dtype + + original_base_model_prefix = kimi_cls.base_model_prefix + kimi_cls.base_model_prefix = "" + try: + model, loading_info = kimi_cls.from_pretrained(str(temp_model_path), **model_kwargs) + finally: + kimi_cls.base_model_prefix = original_base_model_prefix + finally: + if tempdir_context is not None: + tempdir_context.__exit__(None, None, None) + + unexpected_keys = loading_info["unexpected_keys"] + missing_keys = loading_info["missing_keys"] + mismatched_keys = loading_info["mismatched_keys"] + if unexpected_keys or missing_keys or mismatched_keys: + raise RuntimeError( + "Failed to load the stripped Kimi K2.5 checkpoint slice cleanly. " + f"missing={missing_keys}, unexpected={unexpected_keys}, mismatched={mismatched_keys}" + ) + model.eval() + tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) + processor = AutoProcessor.from_pretrained(str(model_path), trust_remote_code=True) + + print(f"Loaded model: {type(model).__name__}") + print(f"Tokenizer vocab size: {tokenizer.vocab_size}") + print(f"Processor type: {type(processor).__name__}") + + return model, tokenizer, processor + + +def load_kimi_k25_layer_subset_model( + *, + model_path: Path | None = None, + num_vision_layers: int = NUM_VISION_LAYERS, + num_text_layers: int = NUM_TEXT_LAYERS, + loaded_expert_ids=LOADED_EXPERT_IDS, + num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, + dtype=torch.float32, + seed: int = 1234, + subset_model_path: Path | None = None, +): + set_deterministic(seed) + resolved_model_path = Path(model_path) if model_path is not None else resolve_model_path() + config = prepare_config(resolved_model_path) + kimi_cls = load_kimi_k25_class(resolved_model_path) + + model, tokenizer, processor = load_layer_subset_model( + model_path=resolved_model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=num_vision_layers, + num_text_layers=num_text_layers, + loaded_expert_ids=loaded_expert_ids, + num_experts_per_tok=num_experts_per_tok, + dtype=dtype, + subset_model_path=subset_model_path, + ) + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +@torch.no_grad() +def run_kimi_k25_hf_model_on_pytorch(model, processor, inputs, max_gen_len): + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print("Original HF Model Outputs (Torch CPU):") + print("Completion:", repr(py_output)) + return output_tokens + + +@torch.no_grad() +def run_kimi_k25_hf_model_on_pytorch_CB(model, processor, images, queries, max_gen_len): + generated_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for idx, (image, query) in enumerate(zip(images, queries)): + conversation = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": query}, + ], + }, + ] + inputs = processor(messages=conversation, add_generation_prompt=True, tokenize=False, return_tensors="pt") + generated_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + for _ in range(max_gen_len): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + output_tokens = torch.cat(new_tokens, dim=1).squeeze(0) + py_output = processor.tokenizer.decode(output_tokens.tolist()).strip() + print(f"Original HF Model Outputs (Torch CPU) for prompt {idx}:") + print("Query:", repr(query)) + print("Completion:", repr(py_output)) + generated_tokens.append(output_tokens.numpy()) + + return generated_tokens