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..23466f53f68 --- /dev/null +++ b/backends/qualcomm/_passes/fuse_consecutive_reshape.py @@ -0,0 +1,34 @@ +# 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): + "For Demo Purpose Only, could possibly be done in other passes." + + 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/backends/qualcomm/quantizer/custom_annotation.py b/backends/qualcomm/quantizer/custom_annotation.py index 54456877c1f..c0ca9756ccb 100644 --- a/backends/qualcomm/quantizer/custom_annotation.py +++ b/backends/qualcomm/quantizer/custom_annotation.py @@ -276,6 +276,273 @@ 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: + """ + 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): + 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 2ed7e49d2a8..7b08851c73d 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 @@ -57,7 +69,8 @@ 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 + # uint8 because annotate 8bit kv io + 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." @@ -106,12 +119,29 @@ 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] + + # 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"] + 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..6164507427c 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,47 @@ 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): + """ + 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) + 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): + 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 +61,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, } @@ -115,16 +151,37 @@ def _qnn_attention_mask( ): 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 = 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): + 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) + with torch.no_grad(): + 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) + + 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__() @@ -134,32 +191,81 @@ 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.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() + 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 + # 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 + ) + logging.info( + f"Replaced in-graph RoPE with precomputed tables (max_seq_len={self.max_seq_len})." ) - self._register_attention_mask_for_4_53(self.exportable_module) - 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) + 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) + 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: torch.Tensor, cache_position: torch.Tensor): - return self.exportable_module( - input_ids=input_ids, cache_position=cache_position + 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..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,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, @@ -43,6 +44,9 @@ FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) +# Use this method name so we can reuse qnn_llama_runner +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", @@ -126,13 +130,13 @@ def source_transform( return self def _tag_ios(self, node, fixed_point_type, config): - # shape of k caches and v caches + 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 +151,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 +162,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 +184,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), -65535.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], dtype=torch.int32), + 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()) @@ -253,6 +306,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, @@ -297,11 +351,12 @@ 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( - 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 +365,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 +376,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 +389,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()], )