From b938412f40c12dac9fb20f5f5a439f4b1eaff8f1 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Thu, 23 Jul 2026 16:33:56 +0800 Subject: [PATCH 1/7] Initial done --- examples/qualcomm/oss_scripts/hf_causal_lm.py | 52 +++++++-- .../llm_utils/decoder_model_wrapper.py | 109 +++++++++++++++--- .../llm_utils/qnn_decoder_model_manager.py | 92 ++++++++++++--- 3 files changed, 209 insertions(+), 44 deletions(-) diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index 2ed7e49d2a8..a381d38081c 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -35,6 +35,18 @@ PTE_FILENAME = "hf_causal_lm_qnn" +# Map the HF decoder_model keys (HUGGING_FACE_REPO_IDS) to the version strings +# the shared qnn_llama_runner understands (see runner.cpp Runner()). +DECODER_MODEL_VERSION = { + "llama3_2-1b": "llama3", + "qwen2_5-0_5b": "qwen2_5", + "qwen2_5-1_5b_instruct": "qwen2_5", + "qwen2_5-0_5b_instruct": "qwen2_5", + "qwen3-0_6b": "qwen3", + "smollm2_135m": "smollm2_135m", + "granite-3_3-2b": "granite", +} + def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 @@ -106,12 +118,30 @@ def inference(args: argparse, qnn_config: QnnConfig): def post_process(): with open(f"{args.artifact}/outputs/result.txt", "r") as f: - outputs.append(f.read()) + text = f.read() + # In tokenized-prompt mode the runner echoes the prompt-file path instead + # of the prompt text; drop it and prepend the real prompt for readability. + prefix = os.path.basename(tokenized_prompt_path) + if text.startswith(prefix): + text = text[len(prefix) :] + outputs.append(args.prompt + text) model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1] seq_len = args.max_seq_len + runner_bin = "examples/qualcomm/oss_scripts/llama/qnn_llama_runner" + decoder_model_version = DECODER_MODEL_VERSION[args.decoder_model] + + # The base (non-instruct) HF models were not trained on the runner's chat + # template. Tokenize the raw prompt here (matching the Python calibration + # path) and feed it via --tokenized_prompt so the runner skips + # get_formatted_prompt. File format: raw little-endian uint64 tokens. + import numpy as np + + prompt_token_ids = tokenizer(args.prompt)["input_ids"] + tokenized_prompt_path = f"{args.artifact}/tokenized_prompt.raw" + np.asarray(prompt_token_ids, dtype=np.uint64).tofile(tokenized_prompt_path) if args.enable_x86_64: # x86 emulator is intended for CI and not performance. Check only the first few tokens. seq_len = min(seq_len, 16) @@ -121,13 +151,15 @@ def post_process(): runner_cmd = " ".join( [ f"export LD_LIBRARY_PATH={qnn_sdk}/lib/{target}/:{args.build_folder}/lib &&", - f"{args.build_folder}/examples/models/llama/llama_main", - f'--prompt "{args.prompt}"', + f"{args.build_folder}/{runner_bin}", + f"--tokenized_prompt {tokenized_prompt_path}", + f"--decoder_model_version {decoder_model_version}", + "--eval_mode 0", f"--tokenizer_path {tokenizer_json_path}", f"--model_path {pte_path}", f"--seq_len {seq_len}", "--temperature 0", - f" > {output_data_folder}/result.txt", + f"--output_path {output_data_folder}/result.txt", ] ) subprocess.run( @@ -141,23 +173,25 @@ def post_process(): runner_cmd = " ".join( [ f"cd {workspace} &&", - "./llama_main", - f'--prompt "{args.prompt}"', + "./qnn_llama_runner", + "--tokenized_prompt tokenized_prompt.raw", + f"--decoder_model_version {decoder_model_version}", + "--eval_mode 0", "--tokenizer_path tokenizer.json", f"--model_path {PTE_FILENAME}.pte", f"--seq_len {seq_len}", "--temperature 0", - " > outputs/result.txt", + "--output_path outputs/result.txt", ] ) adb = SimpleADB( qnn_config=qnn_config, pte_path=pte_path, workspace=workspace, - runner="examples/models/llama/llama_main", + runner=runner_bin, ) # No pregen inputs, input_list is not required - adb.push(inputs=[], files=[tokenizer_json_path]) + adb.push(inputs=[], files=[tokenizer_json_path, tokenized_prompt_path]) adb.execute(custom_runner_cmd=runner_cmd) adb.pull(host_output_path=args.artifact, callback=post_process) diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py index 8dc334baf28..7eb1c240640 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -12,12 +12,50 @@ import torch import transformers from transformers import GenerationConfig, PretrainedConfig - -from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM +from transformers.cache_utils import Cache, StaticLayer TRANSFORMERS_VERSION = "4.53.1" +class QnnCustomStaticLayer(StaticLayer): + """StaticLayer that returns cat(past, new) to attention (so the current layer + sees the full context, matching static_llama.py:494), and stashes the pre-cat + new K/V slice so the wrapper can return only the new slot as graph output.""" + + def __init__(self, max_cache_len): + super().__init__(max_cache_len=max_cache_len) + self.new_keys = None + self.new_values = None + + def update(self, key_states, value_states, cache_kwargs=None): + self.new_keys = key_states + self.new_values = value_states + keys = torch.cat([self.keys, key_states], dim=-2) + values = torch.cat([self.values, value_states], dim=-2) + return keys, values + + def get_mask_sizes(self, cache_position): + return self.max_cache_len, 0 + + +class QnnCustomStaticCache(Cache): + """StaticCache-shaped cache seeded from external past K/V tensors, one pair + per layer. `max_cache_len` is the full context length (past_len + ar_len).""" + + def __init__(self, past_k_list, past_v_list, max_cache_len): + layers = [] + for pk, pv in zip(past_k_list, past_v_list): + layer = QnnCustomStaticLayer(max_cache_len=max_cache_len) + layer.max_batch_size, layer.num_heads, _, layer.head_dim = pk.shape + layer.dtype = pk.dtype + layer.device = pk.device + layer.keys = pk + layer.values = pv + layer.is_initialized = True + layers.append(layer) + super().__init__(layers=layers) + + def save_config_to_constant_methods( config: PretrainedConfig, generation_config: Optional[GenerationConfig] = None, @@ -26,9 +64,10 @@ def save_config_to_constant_methods( # Initialize metadata with values from model config metadata = { "get_bos_id": getattr(config, "bos_token_id", None), - "get_eos_id": getattr(config, "eos_token_id", None), + "get_eos_ids": getattr(config, "eos_token_id", None), "get_vocab_size": getattr(config, "vocab_size", None), "get_max_seq_len": getattr(config, "max_position_embeddings", None), + "get_n_layers": getattr(config, "num_hidden_layers", None), "use_kv_cache": getattr(generation_config, "use_cache", None), "use_sdpa_with_kv_cache": False, } @@ -113,10 +152,12 @@ def _qnn_attention_mask( kv_length: int, **kwargs, ): + # Prefix layout (matches static_llama.py + KVManager): past occupies the + # front of the context [0, n_past), the new token(s) sit in the tail region + # starting at (kv_length - ar_len). This fallback is only used if no explicit + # 4D mask is provided; the QNN flow feeds an explicit mask instead. kv_arange = torch.arange(kv_length, device=cache_position.device) reshaped_cache_position = cache_position.view(-1, 1) - - # Simplest and most efficient way to obtain a causal mask causal_mask = kv_arange <= reshaped_cache_position atten_mask = torch.full((causal_mask.shape[0], kv_length), -65504.0) atten_mask = atten_mask.masked_fill(causal_mask, 0) @@ -134,32 +175,62 @@ def __init__(self, model): model.config, model.generation_config ) logging.info(f"Metadata to be recorded in PTE: {self._metadata}") - self.exportable_module = TorchExportableModuleForDecoderOnlyLM( - self.model, - batch_size=1, - max_cache_len=self._metadata.get("get_max_seq_len"), + + self.num_layers = self.config.num_hidden_layers + self.num_kv_heads = getattr( + self.config, "num_key_value_heads", self.config.num_attention_heads ) - self._register_attention_mask_for_4_53(self.exportable_module) + self.head_dim = self.config.head_dim + self.max_seq_len = self.config.max_seq_len + self.ar_len = self.config.ar_len + self.past_len = self.max_seq_len - self.ar_len + + self._register_attention_mask_for_4_53() - def _register_attention_mask_for_4_53(self, exportable_module: torch.nn.Module): + def _register_attention_mask_for_4_53(self): if transformers.__version__ >= TRANSFORMERS_VERSION: from transformers.masking_utils import AttentionMaskInterface from transformers.modeling_utils import AttentionInterface AttentionInterface.register("qnn_attention", _qnn_attention) AttentionMaskInterface.register("qnn_attention", _qnn_attention_mask) - exportable_module.model.model.config._attn_implementation = "qnn_attention" + self.model.config._attn_implementation = "qnn_attention" self._metadata.update({"use_sdpa_with_kv_cache": False}) def get_example_inputs(self): - example_input_ids = torch.tensor([[1]], dtype=torch.long) - example_cache_position = torch.tensor([0], dtype=torch.long) - return (example_input_ids, example_cache_position) - - def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor): - return self.exportable_module( - input_ids=input_ids, cache_position=cache_position + input_ids = torch.tensor([[1]] * self.ar_len, dtype=torch.long).view(1, -1) + # Explicit additive causal mask, matching static_llama / KVManager: + # 0.0 == attend, large-negative == masked. Shape [B, 1, ar_len, context_len]. + atten_mask = torch.zeros(1, 1, self.ar_len, self.max_seq_len) + input_pos = torch.zeros(1, self.ar_len, dtype=torch.int32) + # K cache is transposed (seq last) to match static_llama: + # K: [B, H, head_dim, past_len] V: [B, H, past_len, head_dim] + past_k = [ + torch.zeros(1, self.num_kv_heads, self.head_dim, self.past_len) + for _ in range(self.num_layers) + ] + past_v = [ + torch.zeros(1, self.num_kv_heads, self.past_len, self.head_dim) + for _ in range(self.num_layers) + ] + return (input_ids, atten_mask, input_pos, past_k, past_v) + + def forward(self, input_ids, atten_mask, input_pos, past_k, past_v): + # Undo the static_llama K transpose so HF sees [B, H, past_len, head_dim]. + past_k_hf = [k.transpose(-1, -2) for k in past_k] + cache = QnnCustomStaticCache(past_k_hf, past_v, max_cache_len=self.max_seq_len) + cache_position = input_pos.view(-1).to(torch.long) + outs = self.model( + input_ids=input_ids, + cache_position=cache_position, + attention_mask=atten_mask, + past_key_values=cache, + use_cache=True, ) + # Return only the new slice, transposing K back to static_llama layout. + new_k = [layer.new_keys.transpose(-1, -2) for layer in cache.layers] + new_v = [layer.new_values for layer in cache.layers] + return outs.logits, new_k, new_v def get_metadata(self): return self._metadata diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index da1ff8e0aba..98f6a836f4a 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -43,6 +43,10 @@ FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) +# Method name the shared qnn_llama_runner expects for the KV path +# (see examples/qualcomm/oss_scripts/llama/runner/runner.cpp). +KV_FORWARD = "kv_forward" + HUGGING_FACE_REPO_IDS = { "llama3_2-1b": "NousResearch/Llama-3.2-1B", "qwen2_5-0_5b": "Qwen/Qwen2.5-0.5B", @@ -69,6 +73,7 @@ def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=Tr config.max_batch_size = batch_size config.enable_spinquant_r3 = enable_spinquant_r3 config.use_cache = True + # config.num_hidden_layers = 1 # Some config has head_dim provided that is different from equation below(e.g., qwen3) if not hasattr(config, "head_dim"): @@ -126,13 +131,16 @@ def source_transform( return self def _tag_ios(self, node, fixed_point_type, config): - # shape of k caches and v caches + # static_llama layout: K is transposed (seq last), V is seq-major. + # K in: [B, H, head_dim, past_len] K out: [B, H, head_dim, ar_len] + # V in: [B, H, past_len, head_dim] V out: [B, H, ar_len, head_dim] + past_len = config.max_seq_len - config.ar_len kv_cache_shape = { - # single head, kv input - (config.head_dim, config.max_seq_len), - (config.max_seq_len, config.head_dim), - # single head, kv output + # K (head_dim, seq) + (config.head_dim, past_len), (config.head_dim, config.ar_len), + # V (seq, head_dim) + (past_len, config.head_dim), (config.ar_len, config.head_dim), } @@ -147,10 +155,7 @@ def _tag_ios(self, node, fixed_point_type, config): quant_io_type = None if node.op == "placeholder": - if ( - len(users := list(node.users)) == 1 - and users[0].meta["val"].size()[-2:] in kv_cache_shape - ): + if node.meta["val"].size()[-2:] in kv_cache_shape: quant_io_type = fixed_point_type["kv_type"] if is_graph_output(node): if node.meta["val"].size()[-2:] in kv_cache_shape: @@ -161,6 +166,8 @@ def _tag_ios(self, node, fixed_point_type, config): return quant_io_type def export(self): + self.graph_module(*self.model_wrapper.get_example_inputs()) + with torch.no_grad(): self.graph_module = torch.export.export( self.graph_module, @@ -181,17 +188,67 @@ def pt2e_calibrate( f"Calibrating with tasks: {calibration_tasks}, limit: {calibration_limit}, calibration_data: {calibration_data}, tokenizer_path: {tokenizer_path}, seq_length: {self.config.max_seq_len}" ) + def _empty_past(): + # K is transposed (seq last) to match static_llama: + # K: [1, H, head_dim, past_len] V: [1, H, past_len, head_dim] + past_k = [ + torch.zeros( + 1, + self.model_wrapper.num_kv_heads, + self.model_wrapper.head_dim, + self.model_wrapper.past_len, + ) + for _ in range(self.model_wrapper.num_layers) + ] + past_v = [ + torch.zeros( + 1, + self.model_wrapper.num_kv_heads, + self.model_wrapper.past_len, + self.model_wrapper.head_dim, + ) + for _ in range(self.model_wrapper.num_layers) + ] + return past_k, past_v + + def _build_mask(n_past, past_len, context_len): + # Prefix layout additive mask [1, 1, ar_len=1, context_len]: + # attend to real past [0, n_past), mask the pad [n_past, past_len), + # attend to the new token at column past_len. + mask = torch.full((1, 1, 1, context_len), -65504.0) + mask[..., :n_past] = 0.0 + mask[..., past_len:] = 0.0 + return mask + def calibrate_template( module: torch.fx.GraphModule, tokenizer, prompts: str, max_len: int ): - # TODO: change criteria & support batch inputs if necessary pos = 0 token_list = tokenizer.encode(prompts, bos=True, eos=False) + past_k, past_v = _empty_past() + past_len = self.model_wrapper.past_len + context_len = self.model_wrapper.max_seq_len + # The prefix buffer holds at most past_len slots, so we can advance + # the position at most past_len times (matching the runner, whose + # seq_len is clamped to context_len). + max_len = min(max_len, past_len) with torch.no_grad(): while token_list[-1] != tokenizer.eos_id and pos < max_len: - cur_pos = torch.tensor([pos], dtype=torch.long) - logits = module(torch.full((1, 1), token_list[pos]), cur_pos) + n_past = min(pos, past_len) + atten_mask = _build_mask(n_past, past_len, context_len) + input_pos = torch.tensor([[n_past]], dtype=torch.int32) + logits, new_k, new_v = module( + torch.full((1, 1), token_list[pos]), + atten_mask, + input_pos, + past_k, + past_v, + ) + # Prefix append: write the new slot into buffer at slot n_past. + for layer in range(self.model_wrapper.num_layers): + past_k[layer][..., :, n_past] = new_k[layer][..., :, 0] + past_v[layer][..., n_past, :] = new_v[layer][..., 0, :] pos += 1 if pos >= len(token_list): token_list.append(torch.argmax(logits, dim=-1).item()) @@ -300,8 +357,8 @@ def to_edge_transform_and_lower_to_qnn( ) with torch.no_grad(): self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( - self.graph_module, - self.model_wrapper.get_example_inputs(), + {KV_FORWARD: self.graph_module}, + {KV_FORWARD: self.model_wrapper.get_example_inputs()}, compiler_spec, constant_methods=self.model_wrapper.get_metadata(), passes_job=self.passes_job, @@ -310,7 +367,9 @@ def to_edge_transform_and_lower_to_qnn( convert_linear_to_conv2d=True, ) - print_delegation_info(self.edge_prog_mgr.exported_program().graph_module) + print_delegation_info( + self.edge_prog_mgr.exported_program(KV_FORWARD).graph_module + ) if not self.use_fp16: logit_out_shape = { ( @@ -319,7 +378,7 @@ def to_edge_transform_and_lower_to_qnn( self.config.vocab_size, ) } - for n in self.edge_prog_mgr.exported_program().graph.nodes: + for n in self.edge_prog_mgr.exported_program(KV_FORWARD).graph.nodes: if n.op == "output": for node, output_encoding in n.meta[QCOM_QUANT_ATTRS_MAP].items(): if node.meta["val"].size() in logit_out_shape: @@ -332,6 +391,7 @@ def to_executorch(self, artifact, pte_filename): executorch_config = ExecutorchBackendConfig( memory_planning_pass=MemoryPlanningPass( alloc_graph_input=False, + alloc_graph_output=False, ), passes=[BuildQuantIo()], ) From 382acf2ad1d1bbb4c972954199f17f01a3e6317d Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Fri, 24 Jul 2026 12:45:34 +0800 Subject: [PATCH 2/7] Use int32 token ids for HF QNN llama flow Align the HF decoder wrapper's token input dtype with static_llama (int32 instead of int64) for both the exported graph signature and the calibration loop. The shared qnn_llama_runner auto-detects the token dtype and narrows accordingly, so no runner change is needed. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py | 2 +- .../qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py index 7eb1c240640..dbaeeabf43c 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -198,7 +198,7 @@ def _register_attention_mask_for_4_53(self): self._metadata.update({"use_sdpa_with_kv_cache": False}) def get_example_inputs(self): - input_ids = torch.tensor([[1]] * self.ar_len, dtype=torch.long).view(1, -1) + input_ids = torch.tensor([[1]] * self.ar_len, dtype=torch.int32).view(1, -1) # Explicit additive causal mask, matching static_llama / KVManager: # 0.0 == attend, large-negative == masked. Shape [B, 1, ar_len, context_len]. atten_mask = torch.zeros(1, 1, self.ar_len, self.max_seq_len) diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index 98f6a836f4a..82a553c9812 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -239,7 +239,7 @@ def calibrate_template( atten_mask = _build_mask(n_past, past_len, context_len) input_pos = torch.tensor([[n_past]], dtype=torch.int32) logits, new_k, new_v = module( - torch.full((1, 1), token_list[pos]), + torch.full((1, 1), token_list[pos], dtype=torch.int32), atten_mask, input_pos, past_k, From a88dac0ced9076bb2f8b100c086bc6c736931838 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Tue, 28 Jul 2026 09:51:04 +0800 Subject: [PATCH 3/7] Enable 8-bit KV cache IO for HF QNN llama flow annotate_kv_8bit_hf tags the KV-cache concat and its graph IO 8-bit for the HF decoder flow, matching static_llama's annotate_kv_8bit. The HF wrapper keeps K stored non-transposed with boundary transposes in forward(), so the walk also traverses the past-KV branch backward (to the K/V placeholder) and the new-KV branch forward (to the K/V graph output) through the pass-through transposes, so both the input and output boundaries are tagged 8-bit. Wire it into pt2e_quantize and set kv_type=uint8 for the 16a*w configs. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../qualcomm/quantizer/custom_annotation.py | 281 ++++++++++++++++++ examples/qualcomm/oss_scripts/hf_causal_lm.py | 5 +- .../llm_utils/qnn_decoder_model_manager.py | 2 + 3 files changed, 287 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/quantizer/custom_annotation.py b/backends/qualcomm/quantizer/custom_annotation.py index 54456877c1f..eb3f9fc53da 100644 --- a/backends/qualcomm/quantizer/custom_annotation.py +++ b/backends/qualcomm/quantizer/custom_annotation.py @@ -276,6 +276,287 @@ def annotate_matmul_input1(node: Node, is_qat: str): annotate_matmul_input1(node.args[1], is_qat=is_qat) +def annotate_kv_8bit_hf( # noqa: C901 + gm: torch.fx.GraphModule, + is_qat=False, +) -> None: + """ + HuggingFace-decoder variant of annotate_kv_8bit (Q @ K^T @ V, 16a8w). + + Same intent as annotate_kv_8bit -- tag the KV cache IO (the cat of past + + new K/V, and the graph-input/output K/V tensors) 8-bit -- but adapted to the + HF graph, which differs from static_llama in two ways: + + 1. q/k/v projections are `aten.linear` (conv2d conversion happens later at + lowering), so the backward walk must stop at `linear` too. + 2. GQA head expansion (`repeat_kv`) sits between the cat and the matmul as + a 5D `expand`/`unsqueeze`/`reshape`. QNN HTP cannot delegate an 8-bit 5D + reshape (validation error 0xc26), which fragments the graph into one + partition per layer. So the walk PASSES THROUGH these repeat_kv ops + without tagging them -- they stay at the default 16-bit and remain + delegatable, while the cat (and thus the KV cache IO) is still 8-bit. + + Net boundary: past_k/v graph inputs, the cat, and new_k/v graph outputs are + 8-bit; the repeat_kv expansion feeding the matmul is 16-bit. + """ + + def annotate_matmul(node: Node, quantization_config: QuantizationConfig): + input_qspec_map = {} + input_act = node.args[0] + input_qspec_map[input_act] = quantization_config.input_activation + input_act1 = node.args[1] + input_qspec_map[input_act1] = quantization_config.weight + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + def annotate_cat(node: Node, quantization_config: QuantizationConfig): + input_nodes = node.args[0] + first_input_node = input_nodes[0] + input_qspec_map = {} + input_qspec_map[first_input_node] = quantization_config.input_activation + share_qparams_with_input_act0_qspec = SharedQuantizationSpec( + (first_input_node, node) + ) + for input_node in input_nodes[1:]: + if input_node not in input_qspec_map: + input_qspec_map[input_node] = share_qparams_with_input_act0_qspec + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=share_qparams_with_input_act0_qspec, + _annotated=True, + ) + + def annotate_single_in_single_out( + node: Node, quantization_config: QuantizationConfig + ) -> None: + input_qspec_map = {} + input_qspec_map[node.args[0]] = quantization_config.input_activation + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + def annotate_single_in_share_out( + node: Node, quantization_config: QuantizationConfig + ) -> None: + input_qspec_map = {} + input_act = node.args[0] + input_qspec_map[input_act] = quantization_config.input_activation + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=SharedQuantizationSpec((input_act, node)), + _annotated=True, + ) + + # repeat_kv (GQA) ops: walk advances through them WITHOUT annotating, so they + # stay 16-bit and QNN can delegate the 5D reshape. + kv_repeat_passthrough = [ + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + ] + + def annotate_past_kv_branch(node: Node, quantization_config: QuantizationConfig): + # Walk the past-KV side of the cat back through pass-through ops (transpose, + # view, reshape, etc.) tagging each 8-bit until we reach the graph + # placeholder (or a non-passthrough op). The final placeholder inherits an + # 8-bit input_activation qspec via annotate_single_in_share_out on the + # closest pass-through op, so BuildQuantIo emits an 8-bit boundary at the + # graph input — matching the QNN graph's expected uint8 KV dtype. + share_out_ops = { + torch.ops.aten.permute.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.transpose.int, + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.flatten.using_ints, + } + for _ in range(16): + if not (isinstance(node, Node) and node.op == "call_function"): + return + if node.target in share_out_ops: + annotate_single_in_share_out(node, quantization_config) + node = node.args[0] + continue + return + + def annotate_new_kv_output_branch( + node: Node, quantization_config: QuantizationConfig + ): + # Walk FORWARD from the cat's new-KV input through pass-through ops that + # feed the graph output, tagging each 8-bit. Mirrors annotate_past_kv_branch + # but on the output boundary. Stops when it hits a non-passthrough user or + # a node with multiple call_function users (branch ambiguity). + share_out_ops = { + torch.ops.aten.permute.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.transpose.int, + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.flatten.using_ints, + } + for _ in range(16): + # Find the single pass-through user that reaches a graph output. + passthrough_users = [ + u + for u in node.users + if u.op == "call_function" and u.target in share_out_ops + ] + if len(passthrough_users) != 1: + return + nxt = passthrough_users[0] + annotate_single_in_share_out(nxt, quantization_config) + node = nxt + + def annotate_matmul_input1(node: Node, is_qat: bool): + if is_qat: + quantization_config_8a8w = get_8a8w_qnn_qat_config( + act_symmetric=True, act_observer=MinMaxObserver + ) + else: + quantization_config_8a8w = get_8a8w_qnn_ptq_config( + act_symmetric=True, act_observer=MinMaxObserver + ) + while isinstance(node, Node) and node.op == "call_function": + if node.target in [ + torch.ops.aten.select.int, + torch.ops.aten.slice.Tensor, + ]: + annotate_single_in_single_out(node, quantization_config_8a8w) + node = node.args[0] + elif node.target in kv_repeat_passthrough: + # Pass through repeat_kv without tagging (stays 16-bit). + node = node.args[0] + elif node.target in [ + torch.ops.aten.permute.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.transpose.int, + torch.ops.aten.flatten.using_ints, + ]: + annotate_single_in_share_out(node, quantization_config_8a8w) + node = node.args[0] + elif node.target == torch.ops.aten.cat.default: + annotate_cat(node, quantization_config_8a8w) + # Also walk the past-KV branch (args[0][0]) through pass-through + # ops to tag them 8-bit, so the graph placeholder inherits the + # 8-bit input_activation qspec and BuildQuantIo emits an 8-bit + # boundary at the K/V graph input. + annotate_past_kv_branch(node.args[0][0], quantization_config_8a8w) + # And walk the new-KV branch FORWARD from the cat input, tagging + # any pass-through ops that lead to the graph output (e.g., the + # boundary transpose in QnnCausalLMExportableModule.forward), so + # the K/V graph output is also 8-bit. + annotate_new_kv_output_branch(node.args[0][1], quantization_config_8a8w) + # Follow the new-kv branch (past kv is args[0][0]). + node = node.args[0][1] + elif node.target in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.sub.Tensor, + torch.ops.aten.matmul.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.linear.default, + ]: + break + else: + print(f"The node ({node}) is not expected in the input1 of the matmul") + node = node.args[0] + + if is_qat: + quantization_config_16a8w = get_16a8w_qnn_qat_config( + act_observer=MinMaxObserver + ) + else: + quantization_config_16a8w = get_16a8w_qnn_ptq_config( + act_observer=MinMaxObserver + ) + + def _is_4d(node: Node) -> bool: + val = node.meta.get("val") + return val is not None and hasattr(val, "dim") and val.dim() == 4 + + def _arg1_reaches_kv_cat(arg1: Node) -> bool: + # Walk arg1 backward following the same ops annotate_matmul_input1 does, + # and return True if we hit a cat.default (the KV-cache concat). This + # picks up KV-cache matmuls even when their arg0 is a 2D constant (e.g. + # R3 spinquant's hadamard rotation matmul(hadamard, cat_K), which fails + # a purely 4D-both-args classifier). + pass_through = { + torch.ops.aten.select.int, + torch.ops.aten.slice.Tensor, + torch.ops.aten.permute.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.transpose.int, + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.flatten.using_ints, + } + node = arg1 + for _ in range(16): # bounded, prevents runaway walks + if not (isinstance(node, Node) and node.op == "call_function"): + return False + if node.target == torch.ops.aten.cat.default: + return True + if node.target in pass_through and node.args: + node = node.args[0] + continue + return False + return False + + # Debug bisect: restrict annotation to K-only, V-only, or both branches. + # Env override: ANNOTATE_KV_HF_BRANCH = "both" | "k" | "v". + import os as _os + + branch_sel = _os.environ.get("ANNOTATE_KV_HF_BRANCH", "both").lower() + + def _branch_of(arg1: Node) -> str: + # HF Q @ K^T matmul: arg[1] = transpose(K...). + # HF attn @ V matmul: arg[1] = reshape/view(V...). + # HF+R3 K-rotation matmul: arg[1] = cat.default (the K cache). + if arg1.op == "call_function": + if arg1.target in ( + torch.ops.aten.transpose.int, + torch.ops.aten.cat.default, + ): + return "k" + if arg1.target in ( + torch.ops.aten.reshape.default, + torch.ops.aten.view.default, + ): + return "v" + return "?" + + for node in gm.graph.nodes: + if not ( + node.op == "call_function" + and node.target == torch.ops.aten.matmul.default + and all(arg.op == "call_function" for arg in node.args) + # Only 4D matmul.default whose arg1 reaches a KV-cache cat are the + # attention matmuls we care about. This filter (arg1 reaches cat) + # replaces the earlier "all args are 4D" filter, which incorrectly + # excluded R3 spinquant's matmul(hadamard, cat_K) whose arg0 is a + # 2D constant. Also excludes HF's 3D RoPE inv_freq matmul. + and _is_4d(node) + and _arg1_reaches_kv_cat(node.args[1]) + ): + continue + br = _branch_of(node.args[1]) + if branch_sel != "both" and br != branch_sel: + continue + annotate_matmul(node, quantization_config_16a8w) + annotate_matmul_input1(node.args[1], is_qat=is_qat) + + def custom_annotate_llama_matmul_16a8w(gm: torch.fx.GraphModule) -> None: # noqa: C901 """ This function is specific for llama matmul op 16a8w. diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index a381d38081c..0b7b976de3b 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -69,7 +69,10 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 "16a16w", ): fixed_point_type["io_type"] = torch.uint16 - fixed_point_type["kv_type"] = torch.uint16 + # KV cache IO is 8-bit: annotate_kv_8bit_hf tags the KV cat 8-bit + # regardless of the activation width (matches static_llama's + # `annotate_kv_8bit` recipe wiring). + fixed_point_type["kv_type"] = torch.uint8 else: raise ValueError( f"No support for quant type {args.ptq}. Support 8a8w, 16a8w, 16a4w and 16a4w_block." diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index 82a553c9812..d1a914d5cec 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -16,6 +16,7 @@ ) from executorch.backends.qualcomm.builders.utils import is_graph_output from executorch.backends.qualcomm.export_utils import make_quantizer +from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit_hf from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.backends.qualcomm.utils.constants import ( QCOM_PASS_ACTIVATE_KEY, @@ -310,6 +311,7 @@ def pt2e_quantize( quantizer = make_quantizer( quant_dtype=quant_dtype, + custom_annotations=(partial(annotate_kv_8bit_hf, is_qat=False),), per_channel_linear=True, per_channel_conv=True, act_observer=MinMaxObserver, From 43cf3bb340a41c62f9629bc7d2ece5d84362f8f9 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 29 Jul 2026 10:15:38 +0800 Subject: [PATCH 4/7] Precompute RoPE tables for the HF QNN llama flow HF computes RoPE inside the traced graph, so the inv_freq math becomes a quantization target. For rope_type=llama3 (llama3.2) inv_freq spans about seven orders of magnitude, and a single per-tensor scale cannot represent it: the low-frequency components that carry long-range position information collapse and the model degenerates into repetitive output under 16a8w. qwen2.5 escapes this only because its default RoPE range is roughly sixteen times tighter. Constant folding cannot help here because the cos/sin computation depends on position_ids, a runtime input. Instead, precompute the tables over the whole position domain [0, max_seq_len) on the host, the way static_llama does, and leave the graph with a gather. The tables are built by calling the original rotary embedding module so any rope_type, scaling and attention_scaling are honored exactly; cos/sin match the original bit-exactly. Only values bounded in [-1, 1] then enter the quantized graph. Removing the per-token inv_freq matmul, cos, sin and their quant/dequant nodes also cuts real work: llama3_2-1b at 16a8w, CL=128 goes from 35.1/31.4 to 49.6/45.9 tokens per second for prefill/decode, and the output is coherent again. fp16 stays coherent. Also align the attention mask fill value with static_llama's PADDING_MASK_VALUE. MHA-to-SHA is left disabled; it still produces a graph HTP rejects. Co-Authored-By: Claude Opus 5 (1M context) --- examples/qualcomm/oss_scripts/hf_causal_lm.py | 2 +- .../llm_utils/decoder_model_wrapper.py | 66 ++++++++++++++++++- .../llm_utils/qnn_decoder_model_manager.py | 12 ++-- 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index 0b7b976de3b..83513dcd837 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -72,7 +72,7 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 # KV cache IO is 8-bit: annotate_kv_8bit_hf tags the KV cat 8-bit # regardless of the activation width (matches static_llama's # `annotate_kv_8bit` recipe wiring). - fixed_point_type["kv_type"] = torch.uint8 + fixed_point_type["kv_type"] = torch.uint16 else: raise ValueError( f"No support for quant type {args.ptq}. Support 8a8w, 16a8w, 16a4w and 16a4w_block." diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py index dbaeeabf43c..d51f089b1ac 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -159,13 +159,59 @@ def _qnn_attention_mask( kv_arange = torch.arange(kv_length, device=cache_position.device) reshaped_cache_position = cache_position.view(-1, 1) causal_mask = kv_arange <= reshaped_cache_position - atten_mask = torch.full((causal_mask.shape[0], kv_length), -65504.0) + # -255 matches static_llama's PADDING_MASK_VALUE (masking_utils.py:12). A + # tighter range keeps the mask input's quant scale small during calibration, + # which preserves precision on the attention-logit space near 0. Using a + # very-large negative (e.g. -65504) forces a coarse quant scale that + # destroys attention numerics under 16a8w — repetitive/degenerate output. + atten_mask = torch.full((causal_mask.shape[0], kv_length), -255.0) atten_mask = atten_mask.masked_fill(causal_mask, 0) atten_mask = atten_mask[None, None, :, :].expand(batch_size, -1, -1, -1) return atten_mask +class QnnPrecomputedRotaryEmbedding(torch.nn.Module): + """Drop-in replacement for HF's ``*RotaryEmbedding`` that serves cos/sin from + precomputed tables instead of computing them in the graph. + + HF computes RoPE inside the traced graph: ``inv_freq @ position_ids`` then + cos/sin. For rope_type="llama3" (llama3.2) ``inv_freq`` spans ~7 orders of + magnitude (min 9.4e-08, max/min 1.06e7), so quantizing those in-graph + intermediates destroys the low-frequency (long-range) position information + and the model degenerates into repetitive output. + + static_llama avoids this by precomputing ``freqs_cos``/``freqs_sin`` on the + host in fp32 and indexing them with ``input_pos`` (static_llama.py:697-698, + 727-731). Only the results -- bounded in [-1, 1] -- enter the graph, which + quantizes cleanly. This module reproduces that: the tables are built once + with the original module's own math (so any rope_type/scaling is honored) + and the graph is left with a gather. + """ + + def __init__(self, rotary_emb: torch.nn.Module, max_seq_len: int, dtype): + super().__init__() + positions = torch.arange(max_seq_len, dtype=torch.long).unsqueeze(0) + # Reuse the original module's forward so rope_type/scaling/attention + # scaling are all applied exactly as HF would. + dummy = torch.zeros(1, max_seq_len, 1, dtype=dtype) + with torch.no_grad(): + cos, sin = rotary_emb(dummy, positions) + self.register_buffer("cos_table", cos[0].to(dtype), persistent=False) + self.register_buffer("sin_table", sin[0].to(dtype), persistent=False) + + def forward(self, x, position_ids): + # position_ids: [batch, seq] -> gather [batch, seq, head_dim] + flat = position_ids.reshape(-1) + cos = self.cos_table.index_select(0, flat).view( + *position_ids.shape, self.cos_table.shape[-1] + ) + sin = self.sin_table.index_select(0, flat).view( + *position_ids.shape, self.sin_table.shape[-1] + ) + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + class QnnCausalLMExportableModule(torch.nn.Module): def __init__(self, model): super().__init__() @@ -186,6 +232,24 @@ def __init__(self, model): self.past_len = self.max_seq_len - self.ar_len self._register_attention_mask_for_4_53() + self._use_precomputed_rope() + + def _use_precomputed_rope(self): + """Swap HF's in-graph RoPE for precomputed cos/sin tables (see + QnnPrecomputedRotaryEmbedding). Keeps HF's plumbing intact -- the decoder + still calls ``self.rotary_emb(hidden_states, position_ids)`` -- but the + wide-dynamic-range inv_freq math now happens on the host in fp32.""" + decoder = self.model.model + rotary_emb = getattr(decoder, "rotary_emb", None) + if rotary_emb is None: + logging.warning("No rotary_emb found; skipping RoPE precompute.") + return + decoder.rotary_emb = QnnPrecomputedRotaryEmbedding( + rotary_emb, self.max_seq_len, self.model.dtype + ) + logging.info( + f"Replaced in-graph RoPE with precomputed tables (max_seq_len={self.max_seq_len})." + ) def _register_attention_mask_for_4_53(self): if transformers.__version__ >= TRANSFORMERS_VERSION: diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index d1a914d5cec..f19a43ce9a7 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -16,7 +16,9 @@ ) from executorch.backends.qualcomm.builders.utils import is_graph_output from executorch.backends.qualcomm.export_utils import make_quantizer -from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit_hf +from executorch.backends.qualcomm.quantizer.custom_annotation import ( # noqa: F401 + annotate_kv_8bit_hf, # kept for the commented-out 8-bit KV IO wiring below +) from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.backends.qualcomm.utils.constants import ( QCOM_PASS_ACTIVATE_KEY, @@ -74,7 +76,6 @@ def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=Tr config.max_batch_size = batch_size config.enable_spinquant_r3 = enable_spinquant_r3 config.use_cache = True - # config.num_hidden_layers = 1 # Some config has head_dim provided that is different from equation below(e.g., qwen3) if not hasattr(config, "head_dim"): @@ -216,7 +217,7 @@ def _build_mask(n_past, past_len, context_len): # Prefix layout additive mask [1, 1, ar_len=1, context_len]: # attend to real past [0, n_past), mask the pad [n_past, past_len), # attend to the new token at column past_len. - mask = torch.full((1, 1, 1, context_len), -65504.0) + mask = torch.full((1, 1, 1, context_len), -65535.0) mask[..., :n_past] = 0.0 mask[..., past_len:] = 0.0 return mask @@ -311,7 +312,7 @@ def pt2e_quantize( quantizer = make_quantizer( quant_dtype=quant_dtype, - custom_annotations=(partial(annotate_kv_8bit_hf, is_qat=False),), + # custom_annotations=(partial(annotate_kv_8bit_hf, is_qat=False),), per_channel_linear=True, per_channel_conv=True, act_observer=MinMaxObserver, @@ -340,6 +341,8 @@ def extract_linear_nodes(graph): tokenizer_path, ) self.graph_module = convert_pt2e(self.graph_module) + # from executorch.backends.qualcomm.utils.utils import draw_graph + # draw_graph("qdq", "./hf_llm", self.graph_module) self.passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True self.passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ @@ -356,6 +359,7 @@ def to_edge_transform_and_lower_to_qnn( compiler_spec = generate_qnn_executorch_compiler_spec( soc_model=get_soc_to_chipset_map()[soc_model], backend_options=backend_options, + # use_mha2sha=True, ) with torch.no_grad(): self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( From 5cad6e3a89853483bb2231035907976b25865c41 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 29 Jul 2026 13:33:18 +0800 Subject: [PATCH 5/7] annotate kv 8bit io --- examples/qualcomm/oss_scripts/hf_causal_lm.py | 2 +- .../oss_scripts/llm_utils/qnn_decoder_model_manager.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index 83513dcd837..0b7b976de3b 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -72,7 +72,7 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 # KV cache IO is 8-bit: annotate_kv_8bit_hf tags the KV cat 8-bit # regardless of the activation width (matches static_llama's # `annotate_kv_8bit` recipe wiring). - fixed_point_type["kv_type"] = torch.uint16 + fixed_point_type["kv_type"] = torch.uint8 else: raise ValueError( f"No support for quant type {args.ptq}. Support 8a8w, 16a8w, 16a4w and 16a4w_block." diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index f19a43ce9a7..b4ae0cb5b46 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -312,7 +312,7 @@ def pt2e_quantize( quantizer = make_quantizer( quant_dtype=quant_dtype, - # custom_annotations=(partial(annotate_kv_8bit_hf, is_qat=False),), + custom_annotations=(partial(annotate_kv_8bit_hf, is_qat=False),), per_channel_linear=True, per_channel_conv=True, act_observer=MinMaxObserver, @@ -359,7 +359,7 @@ def to_edge_transform_and_lower_to_qnn( compiler_spec = generate_qnn_executorch_compiler_spec( soc_model=get_soc_to_chipset_map()[soc_model], backend_options=backend_options, - # use_mha2sha=True, + use_mha2sha=False, ) with torch.no_grad(): self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( From 321a02b16e91555929229bc765a9ae9ad29a09c6 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 29 Jul 2026 20:54:10 +0800 Subject: [PATCH 6/7] Support mha_to_sha --- backends/qualcomm/_passes/__init__.py | 2 + .../_passes/fuse_consecutive_reshape.py | 45 +++++++++++++++++++ backends/qualcomm/_passes/qnn_pass_manager.py | 2 + .../llm_utils/qnn_decoder_model_manager.py | 3 +- 4 files changed, 51 insertions(+), 1 deletion(-) create mode 100644 backends/qualcomm/_passes/fuse_consecutive_reshape.py diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 137a80dcc06..ebcdedbcaf7 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -49,6 +49,7 @@ from .fixed_linear_keep_dim import FixedLinearKeepDim from .fold_qdq import FoldQDQ from .fuse_consecutive_cast import FuseConsecutiveCast +from .fuse_consecutive_reshape import FuseConsecutiveReshape from .fuse_consecutive_transpose import FuseConsecutiveTranspose from .i64_to_i32 import I64toI32 from .insert_cast_for_fp_act_quantized_weight import InsertCastForFpActQuantizedWeight @@ -114,6 +115,7 @@ FixedLinearKeepDim, FoldQDQ, FuseConsecutiveCast, + FuseConsecutiveReshape, FuseConsecutiveTranspose, I64toI32, InsertCastForFpActQuantizedWeight, diff --git a/backends/qualcomm/_passes/fuse_consecutive_reshape.py b/backends/qualcomm/_passes/fuse_consecutive_reshape.py new file mode 100644 index 00000000000..ec27108a037 --- /dev/null +++ b/backends/qualcomm/_passes/fuse_consecutive_reshape.py @@ -0,0 +1,45 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass + + +class FuseConsecutiveReshape(ExportPass): + """ + Collapse chains of consecutive view_copy into a single view_copy. A + view_copy is order-preserving, so ``view(view(x, A), B) == view(x, B)`` for + any valid shapes; only the final shape matters. + + This matters for the ConvInplaceLinear -> ConvertMhaToSha interaction: the + ConvertLinearToConv2d restore emits a rank-3 view_copy that sits between the + conv and the head-making rank-4 view_copy. That intermediate view stops + ConvertMhaToSha's _is_making_mha matcher from reaching the rank-4 reshape, so + the Q/K/V projection weight is never split per head. Fusing the views away + restores the conv -> permute -> view(4D) shape the matcher expects. + """ + + def __init__(self): + super().__init__() + self.view = exir_ops.edge.aten.view_copy.default + + def _fuse(self, graph_module: torch.fx.GraphModule): + for node in graph_module.graph.nodes: + if node.target != self.view: + continue + src = node.args[0] + while isinstance(src, torch.fx.Node) and src.target == self.view: + src = src.args[0] + if src is not node.args[0]: + node.args = (src, *node.args[1:]) + + def call(self, graph_module: torch.fx.GraphModule): + self._fuse(graph_module) + dead_code_elimination_pass(graph_module) + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 82600c90938..8b4118ae629 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -53,6 +53,7 @@ FixedLinearKeepDim, FoldQDQ, FuseConsecutiveCast, + FuseConsecutiveReshape, FuseConsecutiveTranspose, I64toI32, InsertCastForFpActQuantizedWeight, @@ -241,6 +242,7 @@ def get_preprocess_passes( """Return preprocess pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ FoldQDQ, + FuseConsecutiveReshape, ConvertMhaToSha, InsertRequantize, InsertIOQDQ, diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index b4ae0cb5b46..9dd58995bb1 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -76,6 +76,7 @@ def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=Tr config.max_batch_size = batch_size config.enable_spinquant_r3 = enable_spinquant_r3 config.use_cache = True + # config.num_hidden_layers = 1 # Some config has head_dim provided that is different from equation below(e.g., qwen3) if not hasattr(config, "head_dim"): @@ -359,7 +360,7 @@ def to_edge_transform_and_lower_to_qnn( compiler_spec = generate_qnn_executorch_compiler_spec( soc_model=get_soc_to_chipset_map()[soc_model], backend_options=backend_options, - use_mha2sha=False, + use_mha2sha=True, ) with torch.no_grad(): self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( From 6b077342875973fdfd39ff358ce7488dac3ba65a Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 29 Jul 2026 21:47:09 +0800 Subject: [PATCH 7/7] Remove comments --- .../_passes/fuse_consecutive_reshape.py | 13 +----- .../qualcomm/quantizer/custom_annotation.py | 20 ++------- examples/qualcomm/oss_scripts/hf_causal_lm.py | 11 ++--- .../llm_utils/decoder_model_wrapper.py | 43 +++---------------- .../llm_utils/qnn_decoder_model_manager.py | 13 +----- 5 files changed, 17 insertions(+), 83 deletions(-) diff --git a/backends/qualcomm/_passes/fuse_consecutive_reshape.py b/backends/qualcomm/_passes/fuse_consecutive_reshape.py index ec27108a037..23466f53f68 100644 --- a/backends/qualcomm/_passes/fuse_consecutive_reshape.py +++ b/backends/qualcomm/_passes/fuse_consecutive_reshape.py @@ -12,18 +12,7 @@ class FuseConsecutiveReshape(ExportPass): - """ - Collapse chains of consecutive view_copy into a single view_copy. A - view_copy is order-preserving, so ``view(view(x, A), B) == view(x, B)`` for - any valid shapes; only the final shape matters. - - This matters for the ConvInplaceLinear -> ConvertMhaToSha interaction: the - ConvertLinearToConv2d restore emits a rank-3 view_copy that sits between the - conv and the head-making rank-4 view_copy. That intermediate view stops - ConvertMhaToSha's _is_making_mha matcher from reaching the rank-4 reshape, so - the Q/K/V projection weight is never split per head. Fusing the views away - restores the conv -> permute -> view(4D) shape the matcher expects. - """ + "For Demo Purpose Only, could possibly be done in other passes." def __init__(self): super().__init__() diff --git a/backends/qualcomm/quantizer/custom_annotation.py b/backends/qualcomm/quantizer/custom_annotation.py index eb3f9fc53da..c0ca9756ccb 100644 --- a/backends/qualcomm/quantizer/custom_annotation.py +++ b/backends/qualcomm/quantizer/custom_annotation.py @@ -281,23 +281,9 @@ def annotate_kv_8bit_hf( # noqa: C901 is_qat=False, ) -> None: """ - HuggingFace-decoder variant of annotate_kv_8bit (Q @ K^T @ V, 16a8w). - - Same intent as annotate_kv_8bit -- tag the KV cache IO (the cat of past + - new K/V, and the graph-input/output K/V tensors) 8-bit -- but adapted to the - HF graph, which differs from static_llama in two ways: - - 1. q/k/v projections are `aten.linear` (conv2d conversion happens later at - lowering), so the backward walk must stop at `linear` too. - 2. GQA head expansion (`repeat_kv`) sits between the cat and the matmul as - a 5D `expand`/`unsqueeze`/`reshape`. QNN HTP cannot delegate an 8-bit 5D - reshape (validation error 0xc26), which fragments the graph into one - partition per layer. So the walk PASSES THROUGH these repeat_kv ops - without tagging them -- they stay at the default 16-bit and remain - delegatable, while the cat (and thus the KV cache IO) is still 8-bit. - - Net boundary: past_k/v graph inputs, the cat, and new_k/v graph outputs are - 8-bit; the repeat_kv expansion feeding the matmul is 16-bit. + Demo purpose only, could possibly be merged into annotate_kv_8bit in future. + Main idea is to tag io as 8bit here. + Will ensure all ops are properly tagged like annotate_kv_8bit during official PR. """ def annotate_matmul(node: Node, quantization_config: QuantizationConfig): diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index 0b7b976de3b..7b08851c73d 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -69,9 +69,7 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 "16a16w", ): fixed_point_type["io_type"] = torch.uint16 - # KV cache IO is 8-bit: annotate_kv_8bit_hf tags the KV cat 8-bit - # regardless of the activation width (matches static_llama's - # `annotate_kv_8bit` recipe wiring). + # uint8 because annotate 8bit kv io fixed_point_type["kv_type"] = torch.uint8 else: raise ValueError( @@ -136,10 +134,9 @@ def post_process(): runner_bin = "examples/qualcomm/oss_scripts/llama/qnn_llama_runner" decoder_model_version = DECODER_MODEL_VERSION[args.decoder_model] - # The base (non-instruct) HF models were not trained on the runner's chat - # template. Tokenize the raw prompt here (matching the Python calibration - # path) and feed it via --tokenized_prompt so the runner skips - # get_formatted_prompt. File format: raw little-endian uint64 tokens. + # Using tokenized prompt as a workaround. + # qnn_llama_runner expects llama3 equals to instruct one while hf llama3 is not. + # Using tokenized prompt so qnn_llama_runner won't append sytstem prompt on user input prompt. import numpy as np prompt_token_ids = tokenizer(args.prompt)["input_ids"] diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py index d51f089b1ac..6164507427c 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -18,9 +18,9 @@ class QnnCustomStaticLayer(StaticLayer): - """StaticLayer that returns cat(past, new) to attention (so the current layer - sees the full context, matching static_llama.py:494), and stashes the pre-cat - new K/V slice so the wrapper can return only the new slot as graph output.""" + """ + Using Custom Cache so model output kv cache like static_llama. + """ def __init__(self, max_cache_len): super().__init__(max_cache_len=max_cache_len) @@ -39,9 +39,6 @@ def get_mask_sizes(self, cache_position): class QnnCustomStaticCache(Cache): - """StaticCache-shaped cache seeded from external past K/V tensors, one pair - per layer. `max_cache_len` is the full context length (past_len + ar_len).""" - def __init__(self, past_k_list, past_v_list, max_cache_len): layers = [] for pk, pv in zip(past_k_list, past_v_list): @@ -152,18 +149,9 @@ def _qnn_attention_mask( kv_length: int, **kwargs, ): - # Prefix layout (matches static_llama.py + KVManager): past occupies the - # front of the context [0, n_past), the new token(s) sit in the tail region - # starting at (kv_length - ar_len). This fallback is only used if no explicit - # 4D mask is provided; the QNN flow feeds an explicit mask instead. kv_arange = torch.arange(kv_length, device=cache_position.device) reshaped_cache_position = cache_position.view(-1, 1) causal_mask = kv_arange <= reshaped_cache_position - # -255 matches static_llama's PADDING_MASK_VALUE (masking_utils.py:12). A - # tighter range keeps the mask input's quant scale small during calibration, - # which preserves precision on the attention-logit space near 0. Using a - # very-large negative (e.g. -65504) forces a coarse quant scale that - # destroys attention numerics under 16a8w — repetitive/degenerate output. atten_mask = torch.full((causal_mask.shape[0], kv_length), -255.0) atten_mask = atten_mask.masked_fill(causal_mask, 0) atten_mask = atten_mask[None, None, :, :].expand(batch_size, -1, -1, -1) @@ -172,31 +160,13 @@ def _qnn_attention_mask( class QnnPrecomputedRotaryEmbedding(torch.nn.Module): - """Drop-in replacement for HF's ``*RotaryEmbedding`` that serves cos/sin from - precomputed tables instead of computing them in the graph. - - HF computes RoPE inside the traced graph: ``inv_freq @ position_ids`` then - cos/sin. For rope_type="llama3" (llama3.2) ``inv_freq`` spans ~7 orders of - magnitude (min 9.4e-08, max/min 1.06e7), so quantizing those in-graph - intermediates destroys the low-frequency (long-range) position information - and the model degenerates into repetitive output. - - static_llama avoids this by precomputing ``freqs_cos``/``freqs_sin`` on the - host in fp32 and indexing them with ``input_pos`` (static_llama.py:697-698, - 727-731). Only the results -- bounded in [-1, 1] -- enter the graph, which - quantizes cleanly. This module reproduces that: the tables are built once - with the original module's own math (so any rope_type/scaling is honored) - and the graph is left with a gather. - """ - def __init__(self, rotary_emb: torch.nn.Module, max_seq_len: int, dtype): super().__init__() positions = torch.arange(max_seq_len, dtype=torch.long).unsqueeze(0) - # Reuse the original module's forward so rope_type/scaling/attention - # scaling are all applied exactly as HF would. - dummy = torch.zeros(1, max_seq_len, 1, dtype=dtype) with torch.no_grad(): - cos, sin = rotary_emb(dummy, positions) + cos, sin = rotary_emb( + torch.zeros(1, max_seq_len, 1, dtype=dtype), positions + ) self.register_buffer("cos_table", cos[0].to(dtype), persistent=False) self.register_buffer("sin_table", sin[0].to(dtype), persistent=False) @@ -244,6 +214,7 @@ def _use_precomputed_rope(self): if rotary_emb is None: logging.warning("No rotary_emb found; skipping RoPE precompute.") return + # HF API does not have a way to replace, so manually replace it. decoder.rotary_emb = QnnPrecomputedRotaryEmbedding( rotary_emb, self.max_seq_len, self.model.dtype ) diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index 9dd58995bb1..6cb5f880114 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -16,9 +16,7 @@ ) from executorch.backends.qualcomm.builders.utils import is_graph_output from executorch.backends.qualcomm.export_utils import make_quantizer -from executorch.backends.qualcomm.quantizer.custom_annotation import ( # noqa: F401 - annotate_kv_8bit_hf, # kept for the commented-out 8-bit KV IO wiring below -) +from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit_hf from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.backends.qualcomm.utils.constants import ( QCOM_PASS_ACTIVATE_KEY, @@ -46,8 +44,7 @@ FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) -# Method name the shared qnn_llama_runner expects for the KV path -# (see examples/qualcomm/oss_scripts/llama/runner/runner.cpp). +# Use this method name so we can reuse qnn_llama_runner KV_FORWARD = "kv_forward" HUGGING_FACE_REPO_IDS = { @@ -76,7 +73,6 @@ def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=Tr config.max_batch_size = batch_size config.enable_spinquant_r3 = enable_spinquant_r3 config.use_cache = True - # config.num_hidden_layers = 1 # Some config has head_dim provided that is different from equation below(e.g., qwen3) if not hasattr(config, "head_dim"): @@ -134,9 +130,6 @@ def source_transform( return self def _tag_ios(self, node, fixed_point_type, config): - # static_llama layout: K is transposed (seq last), V is seq-major. - # K in: [B, H, head_dim, past_len] K out: [B, H, head_dim, ar_len] - # V in: [B, H, past_len, head_dim] V out: [B, H, ar_len, head_dim] past_len = config.max_seq_len - config.ar_len kv_cache_shape = { # K (head_dim, seq) @@ -342,8 +335,6 @@ def extract_linear_nodes(graph): tokenizer_path, ) self.graph_module = convert_pt2e(self.graph_module) - # from executorch.backends.qualcomm.utils.utils import draw_graph - # draw_graph("qdq", "./hf_llm", self.graph_module) self.passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True self.passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][