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

[yydhYYDH]

Description

This PR fixes activation quantization propagation for LLM NPU export paths.

Binary MUL outputs, such as the gate * up product in SwiGLU, cannot safely derive their quantization range from either input. The product range can be orders of magnitude larger than the operand ranges. Reusing an operand scale causes INT16/UINT16 saturation on QNN and corrupts downstream hidden states, which can result in garbled generation.

This PR partially addresses #4716 by fixing the LLM export portion of the issue, specifically the OmniQuant/SmoothQuant activation-quantization propagation used for QNN/NPU models. It does not claim to resolve Qwen3-VL cases covered by #4716.

The changes:

• Preserve calibrated activation ranges for Binary MUL outputs instead of inferring them from an operand.
• Keep integer index/shape tensors and Gather indices outside activation quantization.
• Apply the corrected propagation rules consistently to OmniQuant and SmoothQuant.

The fix is split into two commits:

11a6c407 [LLM:Bugfix] Protect index tensors during NPU quantization
8682b2bc [LLM:Bugfix] Preserve calibrated MUL activation ranges

For Qwen3-0.6B, the old layer 2 SwiGLU product incorrectly inherited an input scale near 0.0018367. After the fix, the calibrated MUL output and down-projection input both use:

scale = 0.1174984994301709
zero  = -29684

Validation used a full Qwen3-0.6B OmniQuant export with 10 calibration samples, 5 epochs, 4-bit weights, and asymmetric 16-bit activations. All 28 MLP MUL tensors matched their corresponding down-projection input scale and zero point.

QNN offline compilation completed for Snapdragon 8 Elite with soc_id=69, dsp_arch=v79, vtcm_mb=8, and max_history_token=0, producing 30 context binary files. Host CPU, OnePlus 13 CPU, and OnePlus 13 QNN runs all exited with code 0 and generated coherent text.

Reproduction commands

OUT=/tmp/qwen3-mnn

python transformers/llm/export/llmexport.py \
  --path /path/to/Qwen3-0.6B \
  --export mnn \
  --dst_path "$OUT" \
  --mnnconvert /path/to/MNN/build/MNNConvert \
  --quant_bit 4 \
  --quant_block 64 \
  --lm_quant_bit 4 \
  --lm_quant_block 64 \
  --omni \
  --omni_epochs 5 \
  --calib_data /tmp/qwen3_wikitext_10.txt \
  --generate_for_npu \
  --seperate_embed \
  --sym \
  --act_bit 16

python transformers/llm/export/npu/generate_llm_qnn.py \
  --model "$OUT" \
  --soc_id 69 \
  --dsp_arch v79 \
  --vtcm_mb 8 \
  --mnn_path /path/to/MNN/build \
  --chunk_size 128 \
  --max_history_token 0

Test prompt

The same prompt was used for host CPU, OnePlus 13 CPU, and OnePlus 13 QNN:

Explain what a transformer model is in one short sentence.

Test results

Quantization metadata validation:

PASS 28 MLP MUL tensors
layer 2 MUL:       scale=0.1174984994301709, zero=-29684
layer 2 down_proj: scale=0.1174984994301709, zero=-29684

Host MNN CPU generated coherent text:

Okay, the user wants a short explanation of what a transformer model is in one sentence. Let me start by recalling what I know. Transformers are a...

OnePlus 13 CPU result:

output: Okay, the user wants a short explanation of what a transformer model is in one sentence. Transformer models are a type of neural network used to process sequential...

OnePlus 13 QNN result:

output: Okay, the user is asking for a short explanation of what a transformer model is in one sentence. Let me start by recalling what I know. A...

Both mobile runs generated coherent text without the previous garbled-token failure. Performance numbers are included only to identify the validated runs; this PR targets correctness rather than performance.

Module

LLM

Type

• Feature
• Bugfix
• Perf
• Refact
• Style
• Doc
• Test
• Chore

Checklist

• Commit message follows [Module:Type] Description format
• Code compiles without errors
• Tested on relevant platform(s)
• No unrelated format or style changes included

[yydhYYDH]

Additional Qwen3-VL validation has been completed on Snapdragon 8 Elite.

This PR now also fixes the Qwen3-VL QNN offline export and runtime path:

• Read the visual model type from llm_config.json instead of inferring it from the model directory name.
• Declare both image_embeds and deepstack_feature as Qwen3-VL visual graph outputs.
• Generate the correct deepstack_embeds shapes for prefill and decode.
• Handle is_mrope and has_deepstack independently to avoid adding DeepStack inputs to models such as Qwen2.5-VL.
• Slice multimodal RoPE positions and DeepStack features using the current chunk offset during chunked prefill.
• Validate the Qwen3-VL visual module outputs before inference.

Qwen3-VL-2B was exported with 4-bit symmetric weights, asymmetric 16-bit activations, OmniQuant using 10 calibration samples for 5 epochs, and a fixed visual input size of 272×608. The model was compiled for Snapdragon 8 Elite with soc_id=69, dsp_arch=v79, vtcm_mb=8, chunk_size=128, and max_history_token=0, producing 30 LLM binaries and one visual binary.

Reproduction commands:

MODEL_DIR=/path/to/Qwen3-VL-2B-Instruct
OUT_DIR=/path/to/qwen3-vl-2b-mnn
BUILD_DIR=/path/to/MNN/build_qnn
CALIB_DATA=/path/to/wikitext-10.txt

python transformers/llm/export/llmexport.py \
  --path "$MODEL_DIR" \
  --export mnn \
  --dst_path "$OUT_DIR" \
  --mnnconvert "$BUILD_DIR/tools/converter/MNNConvert" \
  --quant_bit 4 \
  --quant_block 64 \
  --lm_quant_bit 4 \
  --lm_quant_block 64 \
  --omni \
  --omni_epochs 5 \
  --calib_data "$CALIB_DATA" \
  --generate_for_npu \
  --seperate_embed \
  --sym \
  --act_bit 16

python transformers/llm/export/npu/generate_llm_qnn.py \
  --model "$OUT_DIR" \
  --soc_id 69 \
  --dsp_arch v79 \
  --vtcm_mb 8 \
  --mnn_path "$BUILD_DIR" \
  --chunk_size 128 \
  --max_history_token 0 \
  --image_sizes 272x608

The test prompt was:

<img><hw>608, 272</hw>/path/on/device/cat.jpg</img>Describe this image in one short sentence.

The four device configurations were run with the same image and prompt:

export LD_LIBRARY_PATH=/path/on/device/runtime/lib

/path/on/device/runtime/bin/llm_demo /path/on/device/model/config_qnn.json /path/on/device/cat_prompt.txt
/path/on/device/runtime/bin/llm_demo /path/on/device/model/config_cpu_llm_qnn_visual.json /path/on/device/cat_prompt.txt
/path/on/device/runtime/bin/llm_demo /path/on/device/model/config_qnn_cpu_visual.json /path/on/device/cat_prompt.txt
/path/on/device/runtime/bin/llm_demo /path/on/device/model/config.json /path/on/device/cat_prompt.txt

The same cat image and prompt were tested with four backend combinations on a OnePlus 13:

Visual backend	LLM backend	Generated text
QNN	QNN	A small, fluffy, light-brown baby goat stands on a pebble in front of a lush green background.
QNN	CPU	A small, orange and white kitten stands on a rock with its ears perked up, looking directly at the camera.
CPU	QNN	A cute, fluffy, young, gingerbread-colored, spotted...
CPU	CPU	A young orange and white kitten stands on a stone with its ears perked and wide eyes looking at the camera, surrounded by lush green grass.

All four combinations completed prefill and decode without graph-shape, graph-execution, or runtime errors. The QNN-visual/CPU-LLM result confirms that the QNN visual graph is working correctly. The remaining visual-language quality loss is localized to the quantized QNN LLM path; the current text-only calibration dataset does not cover visual embedding and DeepStack activation distributions.

[Qxinyu]

Verification report + review

I ported the three commits onto our internal tree (the five touched files were byte-identical to master, so git am -3 applied cleanly) and validated on Snapdragon 8 Gen3 (SM8650, soc_id=57 dsp_arch=v75 vtcm_mb=8) with QAIRT 2.46, doing a base-vs-PR A/B where both sides used the same wikitext-2 calibration set (10 × 2200 chars) and identical export flags.

Reproduced, and stronger than the PR description claims

1) Index protection is not only about garbled text — without it the offline graph does not compile. On base, qnn-context-binary-generator fails:

[ERROR] <E> Unsupported input/output datatypes requested for the HTP Op 'Reshape'
        in the node '_rotary_Reshape_output_0_Reshape'
[ERROR] <E> Requested I/O datatype set:
   in[0] : QNN_DATATYPE_INT_32
   out[0]: QNN_DATATYPE_UFIXED_POINT_16
[ERROR] Failed in composeGraphs()  →  RuntimeError: npu_convert.py failed

Base leaked 613 index/shape tensors into activation quantization (2 int Input incl. position_ids, 257 Shape/Rank/Size, 238 Gather indices, 116 Reshape shape-inputs); with the PR it is 0, and compilation succeeds (30 context binaries).

2) The MUL numbers reproduce. Base: 28/28 MLP MUL outputs inherit an operand scale (layer 2 = 0.0018366381, bit-exact with operand[1]). PR: 0/28; layer 2 = scale 0.11750022 / zero -29691 (vs your 0.1174985 / -29684 — the tiny delta is just the different calibration set, which is itself good news: the value is a property of the model, not of the calib data). Quantized-tensor count 3256 → 1408, propagation passes 6 → 3.

3) Qwen3-VL offline path. llm.mnn and visual.mnn both compile with 0 errors; generated IO json carries deepstack_embeds [3,128,2048] / [3,1,2048] (base emits the placeholder [3,1,1]). On device I ran a 3-way A/B with vision on QNN and the LLM on CPU:

build	visual graph outputs	result
base lib + base graph	1 (image_embeds)	runs, but no deepstack (wording degrades)
PR lib + PR graph	2 (+deepstack_feature)	runs, deepstack active, correct answer
PR lib + base graph	1	new guard fires: requires image_embeds and deepstack_feature, but got 1 outputs → image dropped

Also worth noting: base's makeVLIOJson decides the model family from the directory name ("Qwen3" in args.model), so it raised ValueError: Unsupported visual model for a dir I named vl_base. Reading model_type from llm_config.json is clearly the right fix.

---

Blocking issue 1 — block.float() is missing in optimize_weights()

omni_quantizer.py:724 adds block.float() to the "Collecting Feature Map Info" loop, but the per-block loop in optimize_weights() (line 531) runs before it and has no dtype normalization. On a bf16 checkpoint (Qwen3-0.6B and Qwen3-VL-2B are both torch_dtype: bfloat16) with transformers 5.15 / torch 2.13, --omni dies before ever reaching your fix:

RuntimeError: expected m1 and m2 to have the same dtype, but got: float != c10::BFloat16
  omni_quantizer.py optimize_weights -> _safe_forward -> F.linear

I had to add the same one-line shim to the optimize_weights() loop (on both base and PR, to keep the A/B fair) before either side could export. Worth confirming whether your environment loads the checkpoint in fp32 — that would explain why you did not hit it.

Secondary: block.float() mutates the module in place and never restores the dtype, so bf16 → fp32 doubles resident memory for the remainder of the export. Consider restoring afterwards, or using autocast.

Blocking issue 2 — if binary_type == 'MUL': continue is too blunt

omni_quantizer.py:945 / smooth_quantizer.py:768 skip propagation for every MUL. Measured effect on Qwen3-0.6B:

	MUL ops with quant params	quantized tensors
base	145 / 145 (all inherited, i.e. wrong)	3256
PR	28 / 145	1408

Of the 117 that lose their parameters, 112 live in self_attn — inputs q_norm/Mul_1_output_0 × Gather_2_output_0, i.e. the RoPE q*cos / k*sin multiplies, 4 per layer × 28 layers. Those are genuine hot-path activations, and leaving them unparameterized means unquantized tensors sandwiched between quantized ops on HTP (the usual source of extra quantize/dequantize pairs). The PR targets correctness, which is fair, but quantization coverage drops by 57% and the runtime cost is unmeasured.

Suggestion: keep the calibrated value when present (current behaviour), and when it is absent derive the range from the operands instead of dropping it. For symmetric int16, out_scale ≈ sa * sb * 32767. Measured against the 28 calibrated MLP MULs:

derived / calibrated   median 1.79x   range 0.52 – 16.3x   ← errs high = loses precision, never saturates
base    / calibrated   median 0.313x  range 0.0156 – 1.65x ← errs low  = saturates (layer 2 is 64x off)

So the derived rule is a safe fallback and strictly better than both inheriting (base) and dropping (current PR) for tensors the calibrator never observed.

Related: REALDIV (57 ops) has the same range-expansion property and is not handled. They happen to lose their parameters via the new protected propagation, but that is luck rather than design.

Blocking issue 3 — gen_position_ids relies on an out-of-range fallback

omni.cpp:1782, axisValue(axis, mContext->all_seq_len + i), is the right direction: with chunked prefill the old axisValue(axis, i) indexes the head of mPositionIds for every chunk, so chunks ≥ 2 get the first chunk's H/W. I confirmed the runtime effect by forcing chunked prefill on CPU ("chunk_limits": [128, 1]) with Qwen3-VL-2B: base fails hard, PR does not —

base: Broad cast error, dim1 = 265, dim2 = 128
      Compute Shape Error for /Add_output_0
      onForward returned no outputs. seqLen=128, inDecode=0  →  LLM in error state
PR  : runs to completion

(so the forwardRaw deepstack alignment fixes a real hard failure — 265 is the full-prompt deepstack length vs the 128 chunk.)

However, once all_seq_len + i exceeds mPositionIds' length, the lambda's return i fallback kicks in and mRoPE silently degenerates to 1-D sequential positions. Instrumented dump from that run:

[POSDBG] all_seq_len=0   seq_len=128 listSize=273 | ax0[...last=4] ax1[...last=11] ax2[...last=15]
[POSDBG] all_seq_len=256 seq_len=128 listSize=273 | ax0[...last=383] ax1[...last=383] ax2[...last=383]
                                          ^ out of range        ^ all three axes equal = spatial structure lost

The overflow region is masked out so it appears harmless, but depending on an accidental "fallback returns the index itself" is fragile — an explicit clamp (or assert) would be better. Also note the hunyuan branch keeps the old formula, so the same chunked-prefill bug remains there and the two paths now diverge semantically.

Happy to re-run anything above; the exports, offline binaries and instrumented builds are still on hand.

[yydhYYDH]

Thank you for the detailed review and thorough reproduction! I will look into these blocking issues and get them resolved.