Skip to content

[QNN:Bugfix] Preserve calibrated MUL activation ranges for QNN export #4765

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
yydhYYDH wants to merge 3 commits into
base: master
Choose a base branch
from
Open

[QNN:Bugfix] Preserve calibrated MUL activation ranges for QNN export #4765

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3c49e4e
[QNN:Bugfix] Protect index tensors during NPU quantization
yydhYYDH Aug 13, 2026
454c380
[QNN:Bugfix] Preserve calibrated MUL activation ranges
yydhYYDH Aug 13, 2026
258edd3
[QNN:Bugfix] Fix Qwen3-VL offline graph inputs
yydhYYDH Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 44 additions & 3 deletions transformers/llm/engine/src/omni.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -825,6 +825,15 @@ std::vector<int> Omni::qwen2VisionProcess(VARP image) {
MNN::Express::Variable::save({patches, position_ids, attention_mask}, "input.mnn");
#endif
auto outputs = mVisionModule->onForward(moduleInputs);
if (outputs.empty() || outputs[0] == nullptr || outputs[0]->getInfo() == nullptr) {
MNN_ERROR("Qwen VL vision module returned no valid image embedding.\n");
return {};
}
if (isQwen3VL && outputs.size() != 2) {
MNN_ERROR("Qwen3-VL vision module requires image_embeds and deepstack_feature, but got %zu outputs.\n",
outputs.size());
return {};
}
auto imageEmbedding = outputs[0];
if (outputs.size() == 2) {
mDeepStackEmbeddings.push_back(outputs[1]);
Expand Down Expand Up @@ -1766,8 +1775,12 @@ VARP Omni::gen_position_ids(int seq_len) {
};
for (int i = 0; i < seq_len; i++) {
for (int axis = 0; axis < axes; axis++) {
int offset = (hunyuan && axis > 0) ? 0 : mContext->all_seq_len;
ptr[i + seq_len * axis] = axisValue(axis, i) + offset;
if (hunyuan) {
int offset = axis > 0 ? 0 : mContext->all_seq_len;
ptr[i + seq_len * axis] = axisValue(axis, i) + offset;
} else {
ptr[i + seq_len * axis] = axisValue(axis, mContext->all_seq_len + i);
}
}
}
if (mTalker) {
Expand All @@ -1785,7 +1798,35 @@ VARP Omni::gen_position_ids(int seq_len) {

std::vector<Express::VARP> Omni::forwardRaw(Express::VARP hiddenState, Express::VARP mask, Express::VARP inputPos, Express::VARPS extraArgs) {
MNN::Express::ExecutorScope s(mExecutor);
extraArgs.insert(extraArgs.end(), mExtraArgs.begin(), mExtraArgs.end());
if (mConfig->has_deepstack() && mExtraArgs.size() == 1) {
auto deepstack = mExtraArgs[0];
auto deepstackInfo = deepstack->getInfo();
auto hiddenInfo = hiddenState->getInfo();
if (deepstackInfo != nullptr && hiddenInfo != nullptr && deepstackInfo->dim.size() == 3) {
const int targetLength = hiddenInfo->dim[mSeqLenIndex];
const int sourceLength = deepstackInfo->dim[1];
const int hiddenSize = deepstackInfo->dim[2];
const int sourceOffset = mContext->gen_seq_len > 0 ? sourceLength : mContext->all_seq_len;
if (sourceOffset != 0 || sourceLength != targetLength) {
auto alignedDeepstack = Express::_Input({deepstackInfo->dim[0], targetLength, hiddenSize}, NCHW);
auto dst = alignedDeepstack->writeMap<float>();
const auto src = deepstack->readMap<float>();
::memset(dst, 0, alignedDeepstack->getInfo()->size * sizeof(float));
const int copyLength = std::max(0, std::min(targetLength, sourceLength - sourceOffset));
if (src != nullptr && copyLength > 0) {
for (int i = 0; i < deepstackInfo->dim[0]; ++i) {
::memcpy(dst + i * targetLength * hiddenSize,
src + (i * sourceLength + sourceOffset) * hiddenSize,
copyLength * hiddenSize * sizeof(float));
}
}
deepstack = alignedDeepstack;
}
}
extraArgs.emplace_back(deepstack);
} else {
extraArgs.insert(extraArgs.end(), mExtraArgs.begin(), mExtraArgs.end());
}
if (mIsEmbedding) {
std::vector<VARP> inputs{hiddenState, mask, inputPos};
if (!extraArgs.empty()) {
Expand Down
6 changes: 5 additions & 1 deletion transformers/llm/export/llmexport.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,7 +396,11 @@ def export_config(self, mnn_config = False):
config['dit_solver'] = 1
if self.model_type == "gemma3":
config.update({'precision': "normal"})
if (hasattr(self, 'visual') and self.visual is not None) or (hasattr(self, 'visual') and self.audio is not None):
is_visual = hasattr(self, 'visual') and self.visual is not None
is_audio = hasattr(self, 'audio') and self.audio is not None
if is_visual or is_audio:
config['is_visual'] = is_visual
config['is_audio'] = is_audio
config['mllm'] = {
'backend_type': "cpu",
"thread_num": 4,
Expand Down
45 changes: 23 additions & 22 deletions transformers/llm/export/npu/generate_llm_qnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def is_embedding_model(config_data, model_dir):
return True
return False

def makeIOJson(args, seq_len, hidden_size, mask_type, is_embedding=False):
def makeIOJson(args, seq_len, hidden_size, mask_type, is_embedding=False, is_mrope=False, has_deepstack=False):
def model_inputs(current_seq_len, logits_index=None):
inputs = [
{
Expand Down Expand Up @@ -104,37 +104,34 @@ def model_inputs(current_seq_len, logits_index=None):
inp["shape"] = [2, 1, 1, 1, 3]
if inp["name"] == "position_ids":
inp["shape"] = [3, 1]
if not is_embedding and "Qwen" in args.model and "VL" in args.model:
if not is_embedding and is_mrope:
cfg = config["configs"]
inputs = cfg[0]["inputs"]
for inp in inputs:
if inp["name"] == "position_ids":
inp["shape"] = [3, seq_len]

new_input = {
"name": "deepstack_embeds",
"shape": [3, 1, 1]
}
inputs.append(new_input)

inputs = cfg[1]["inputs"]
for inp in inputs:
if inp["name"] == "position_ids":
inp["shape"] = [3, 1]

new_input = {
if not is_embedding and has_deepstack:
config["configs"][0]["inputs"].append({
"name": "deepstack_embeds",
"shape": [3, 1, 1]
}
inputs.append(new_input)
"shape": [3, seq_len, hidden_size]
})
config["configs"][1]["inputs"].append({
"name": "deepstack_embeds",
"shape": [3, 1, hidden_size]
})
cache = os.path.join(os.getcwd(), args.cache_path)
with open(os.path.join(cache, 'input.json'), 'w') as f:
f.write(json.dumps(config, indent=4))

def makeVLIOJson(args, image_sizes):
def makeVLIOJson(args, image_sizes, model_type):
configs = []
for w, h in image_sizes:
if "Qwen2.5" in args.model and "VL" in args.model:
if model_type == "qwen2_5_vl":
align_size = 28
grid_h = (round(h / align_size) * align_size) // 14
grid_w = (round(w / align_size) * align_size) // 14
Expand All @@ -148,7 +145,7 @@ def makeVLIOJson(args, image_sizes):
],
"outputs": ["image_embeds"]
}
elif "Qwen3" in args.model or "Qwen3.5" in args.model:
elif model_type in {"qwen3_vl", "qwen3_vl_moe", "qwen3_5_vl"}:
align_size = 32
grid_h = (round(h / align_size) * align_size) // 16
grid_w = (round(w / align_size) * align_size) // 16
Expand All @@ -161,17 +158,17 @@ def makeVLIOJson(args, image_sizes):
{"name": "idx_tensor", "shape": [4, seq_len]},
{"name": "weight_tensor", "shape": [4, seq_len]}
],
"outputs": ["image_embeds"]
"outputs": ["image_embeds", "deepstack_feature"]
}
elif "FastVLM" in args.model:
elif model_type in {"fastvlm", "llava_qwen2"}:
config = {
"inputs": [
{"name": "input_images", "shape": [1, 3, h, w]}
],
"outputs": ["image_embeds"]
}
else:
raise ValueError(f"Unsupported visual model: {args.model}")
raise ValueError(f"Unsupported visual model type: {model_type}")
configs.append(config)

full_config = {"configs": configs}
Expand Down Expand Up @@ -347,10 +344,12 @@ def convert_visual(args):
if not image_sizes:
print("No valid image sizes provided.")
sys.exit(1)
if "FastVLM" in args.model:
with open(os.path.join(args.model, "llm_config.json"), "r", encoding="utf-8") as f:
model_type = json.load(f)["model_type"]
if model_type in {"fastvlm", "llava_qwen2"}:
convert_fastvlm(args, image_sizes)
else:
makeVLIOJson(args, image_sizes)
makeVLIOJson(args, image_sizes, model_type)
inputjson = os.path.join(cache, 'input.json')
ids = list(range(len(image_sizes)))
convert_qnn(args, 'visual.mnn', inputjson, external_file, ids)
Expand All @@ -371,11 +370,13 @@ def convert_llm(args):
if "attention_mask" in config_data:
mask_type = config_data["attention_mask"]
is_embedding = is_embedding_model(config_data, model_dir)
is_mrope = config_data.get("is_mrope", False)
has_deepstack = config_data.get("has_deepstack", False)

ids = [0, 1]
model_name = 'embedding.mnn' if is_embedding else 'llm.mnn'
external_file = os.path.join(model_dir, model_name + '.weight')
makeIOJson(args, args.chunk_size, hidden_size, mask_type, is_embedding)
makeIOJson(args, args.chunk_size, hidden_size, mask_type, is_embedding, is_mrope, has_deepstack)
inputjson = os.path.join(cache, 'input.json')
convert_qnn(args, model_name, inputjson, external_file, ids)

Expand Down
113 changes: 90 additions & 23 deletions transformers/llm/export/utils/omni_quantizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -719,6 +719,9 @@ def _collect_feature_map_optimized(self):

for idx in tqdm(range(len(self.modules)), desc="Collecting Feature Map Info"):
block = self.modules[idx]
# Calibration math is implemented in float32; normalize BF16
# checkpoints before collecting activation ranges.
block.float()
self.to_device(block, self.best_device)

target_ops = SmoothQuantizer.get_all_leaf_modules(block)
Expand Down Expand Up @@ -812,13 +815,68 @@ def _propagate_quant_info(self, mnn_ops, quant_info_dict):

DATA_SELECT_OPS = ['Gather', 'GatherV2', 'GatherND']

# Keep index/shape tensors out of activation quantization. In
# particular, position_ids is INT32 and must not become the input of a
# QNN Reshape whose output is UFIXED_POINT_16.
protected = set()
integer_types = {
'DT_INT8', 'DT_INT16', 'DT_INT32', 'DT_INT64',
'DT_UINT8', 'DT_UINT16', 'DT_UINT32', 'DT_UINT64',
}
for op in mnn_ops:
op_type = op.get('type', '')
inputs = op.get('inputIndexes', [])
outputs = op.get('outputIndexes', [])
main = op.get('main') or {}
if op_type == 'Input' and main.get('dtype') in integer_types:
protected.update(outputs)
elif op_type == 'Const' and main.get('dataType') in integer_types:
protected.update(outputs)
elif op_type in {'Shape', 'Rank', 'Size'}:
protected.update(outputs)
if op_type == 'Reshape' and len(inputs) > 1:
protected.update(inputs[1:])
elif op_type in DATA_SELECT_OPS and len(inputs) > 1:
protected.update(inputs[1:])

# Propagate the protected marker through shape/index-only arithmetic.
changed_protected = True
while changed_protected:
changed_protected = False
for op in mnn_ops:
op_type = op.get('type', '')
inputs = op.get('inputIndexes', [])
outputs = op.get('outputIndexes', [])
if not inputs or not outputs:
continue
if op_type in PASS_THROUGH_OPS:
# For Reshape/Slice/etc. only the data input determines
# the output type; shape/axis inputs must not contaminate
# the feature-map output.
mark = inputs[0] in protected
elif op_type in {'BinaryOp', 'UnaryOp'}:
mark = any(index in protected for index in inputs)
else:
mark = False
if mark:
for index in outputs:
if index not in protected:
protected.add(index)
changed_protected = True

for index in protected:
quant_info_dict.pop(index, None)

print("Start propagating quantization parameters...")
changed = True
pass_round = 0

max_passes = len(mnn_ops) + 1
while changed:
changed = False
pass_round += 1
if pass_round > max_passes:
raise RuntimeError("Quantization parameter propagation did not converge")
update_count = 0

for op in mnn_ops:
Expand All @@ -832,15 +890,16 @@ def _propagate_quant_info(self, mnn_ops, quant_info_dict):
if op_type in PASS_THROUGH_OPS:
source_info = None
for inp_idx in inputs:
if inp_idx in quant_info_dict:
if inp_idx in quant_info_dict and inp_idx not in protected:
source_info = quant_info_dict[inp_idx]
break

if source_info:
for out_idx in outputs:
if out_idx not in quant_info_dict:
quant_info_dict[out_idx] = copy.deepcopy(source_info)
quant_info_dict[out_idx]['index'] = out_idx # 修正 index
if out_idx not in protected and out_idx not in quant_info_dict:
new_info = copy.deepcopy(source_info)
new_info['index'] = out_idx
quant_info_dict[out_idx] = new_info
changed = True
update_count += 1

Expand All @@ -852,9 +911,10 @@ def _propagate_quant_info(self, mnn_ops, quant_info_dict):

if target_info:
for inp_idx in inputs:
if inp_idx not in quant_info_dict:
quant_info_dict[inp_idx] = copy.deepcopy(target_info)
quant_info_dict[inp_idx]['index'] = inp_idx
if inp_idx not in quant_info_dict and inp_idx not in protected:
new_info = copy.deepcopy(target_info)
new_info['index'] = inp_idx
quant_info_dict[inp_idx] = new_info
changed = True
update_count += 1

Expand All @@ -863,26 +923,31 @@ def _propagate_quant_info(self, mnn_ops, quant_info_dict):
out_idx = outputs[0]

# Forward: Data -> Output
if data_idx in quant_info_dict and out_idx not in quant_info_dict:
quant_info_dict[out_idx] = copy.deepcopy(quant_info_dict[data_idx])
quant_info_dict[out_idx]['index'] = out_idx
changed = True
update_count += 1

# Backward: Output -> Data
if out_idx in quant_info_dict and data_idx not in quant_info_dict:
quant_info_dict[data_idx] = copy.deepcopy(quant_info_dict[out_idx])
quant_info_dict[data_idx]['index'] = data_idx
changed = True
update_count += 1
if data_idx in quant_info_dict and out_idx not in protected:
source_info = quant_info_dict[data_idx]
output_info = quant_info_dict.get(out_idx)
if output_info is None or output_info.get('quantInfo') != source_info.get('quantInfo'):
quant_info_dict[out_idx] = copy.deepcopy(source_info)
quant_info_dict[out_idx]['index'] = out_idx
changed = True
update_count += 1

# Do not propagate Gather output parameters back to the
# data input: QNN requires Gather data/output parameters to
# match, while the indices input is always unquantized.

elif op_type == 'BinaryOp':
out_idx = outputs[0]

if out_idx in quant_info_dict:
binary_type = (op.get('main') or {}).get('opType', '')
# Multiplication changes the value range. Neither operand
# nor the product can safely inherit the other's scale;
# retain only independently calibrated parameters.
if binary_type == 'MUL':
continue
if out_idx in quant_info_dict and out_idx not in protected:
target_info = quant_info_dict[out_idx]
for inp_idx in inputs:
if inp_idx not in quant_info_dict:
if inp_idx not in quant_info_dict and inp_idx not in protected:
quant_info_dict[inp_idx] = copy.deepcopy(target_info)
quant_info_dict[inp_idx]['index'] = inp_idx
changed = True
Expand All @@ -891,8 +956,10 @@ def _propagate_quant_info(self, mnn_ops, quant_info_dict):
else:
scales = []
valid_inputs = []
if out_idx in protected:
continue
for inp_idx in inputs:
if inp_idx in quant_info_dict:
if inp_idx in quant_info_dict and inp_idx not in protected:
scales.append(quant_info_dict[inp_idx]['quantInfo']['scale'])
valid_inputs.append(inp_idx)

Expand Down
Loading