From c998d136e62a113c34b580c2fb3cf40183be653b Mon Sep 17 00:00:00 2001 From: yuyazhua Date: Mon, 22 Jun 2026 14:37:33 +0800 Subject: [PATCH] Qualcomm AI Engine Direct - Gemma4 Text model enablement (Unified Architecture) Summary: - Add Gemma4-E2B decoder to the static-llama flow: YOCO KV-sharing (self / donor / cross-decoder layers), per-layer embeddings (PLE). - Support mixed per-layer head_dim, global head_dim --- backends/qualcomm/tests/test_passes.py | 4 +- backends/qualcomm/tests/test_qnn_delegate.py | 16 + examples/models/gemma4/__init__.py | 3 +- examples/models/gemma4/config/e2b_config.json | 3 + .../models/gemma4/text_decoder/__init__.py | 2 +- .../gemma4/text_decoder/convert_weights.py | 67 + examples/models/llama/evaluate/eager_eval.py | 2 + examples/models/llama/model_args.py | 56 +- examples/qualcomm/oss_scripts/llama/README.md | 7 + .../qualcomm/oss_scripts/llama/__init__.py | 60 +- .../oss_scripts/llama/decoder_constants.py | 1 + .../llama/encoder/encoder_config.py | 4 +- examples/qualcomm/oss_scripts/llama/llama.py | 5 + .../oss_scripts/llama/model/__init__.py | 65 +- .../oss_scripts/llama/model/apply_rope.py | 52 - .../llama/model/encoders/__init__.py | 15 + .../{audio_encoder.py => encoders/audio.py} | 0 .../{vision_encoder.py => encoders/vision.py} | 0 .../oss_scripts/llama/model/static_llama.py | 1088 ----------------- .../llama/model/text_decoder/__init__.py | 72 ++ .../model/text_decoder/attention_sink.py | 246 ++++ .../model/text_decoder/blocks/__init__.py | 48 + .../model/text_decoder/blocks/attention.py | 549 +++++++++ .../text_decoder/blocks/decoder_layer.py | 183 +++ .../{ => text_decoder/blocks}/feed_forward.py | 60 +- .../blocks/norm.py} | 0 .../model/{ => text_decoder}/embedding.py | 0 .../llama/model/text_decoder/rope.py | 186 +++ .../llama/model/text_decoder/static_llama.py | 576 +++++++++ .../oss_scripts/llama/qnn_llama_runner.cpp | 12 + .../oss_scripts/llama/runner/kv_manager.cpp | 168 +-- .../oss_scripts/llama/runner/kv_manager.h | 18 +- .../oss_scripts/llama/runner/runner.cpp | 21 +- .../oss_scripts/llama/runner/runner.h | 1 + .../llama/static_llm_quant_recipe.py | 24 + .../qualcomm/oss_scripts/llama/tokenizer.py | 10 +- .../llama/wrappers/attention_sink_wrappers.py | 2 +- .../llama/wrappers/base_component.py | 2 +- .../llama/wrappers/llm_wrappers.py | 83 +- exir/lowered_backend_module.py | 4 +- 40 files changed, 2440 insertions(+), 1275 deletions(-) delete mode 100644 examples/qualcomm/oss_scripts/llama/model/apply_rope.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/encoders/__init__.py rename examples/qualcomm/oss_scripts/llama/model/{audio_encoder.py => encoders/audio.py} (100%) rename examples/qualcomm/oss_scripts/llama/model/{vision_encoder.py => encoders/vision.py} (100%) delete mode 100755 examples/qualcomm/oss_scripts/llama/model/static_llama.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/__init__.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/attention_sink.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/__init__.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/attention.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/decoder_layer.py rename examples/qualcomm/oss_scripts/llama/model/{ => text_decoder/blocks}/feed_forward.py (67%) rename examples/qualcomm/oss_scripts/llama/model/{layernorm.py => text_decoder/blocks/norm.py} (100%) rename examples/qualcomm/oss_scripts/llama/model/{ => text_decoder}/embedding.py (100%) create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/rope.py create mode 100644 examples/qualcomm/oss_scripts/llama/model/text_decoder/static_llama.py diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 1124b01d613..e34c12946e4 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -162,9 +162,7 @@ def test_mha_to_sha(self): from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import ( CausalAttentionMask, ) - from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( - LlamaAttention, - ) + from executorch.examples.qualcomm.oss_scripts.llama.model import LlamaAttention # Initailize model config args = ModelArgs() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 98c5839c515..51d414f8094 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -8144,6 +8144,14 @@ def setUp(self): hellaswag_acc_norm=None, sqnr=10, ), + "gemma4-e2b": TestExampleLLMScript.LlmSpecs( + SM8650=20, + SM8750=30, + pte_size=4_500_000_000, # 4.5 GB + wikitext_ppl=120, + hellaswag_acc_norm=None, + sqnr=10, + ), "glm-1_5b": TestExampleLLMScript.LlmSpecs( SM8650=42, SM8750=52, @@ -8345,6 +8353,14 @@ def test_static_llm_model(self): # noqa: C901 ] ) + if self.model_name == "gemma4-e2b": + cmds.extend( + [ + "--embedding-quantize", + "4,32", + ] + ) + self.add_default_cmds(cmds) p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) diff --git a/examples/models/gemma4/__init__.py b/examples/models/gemma4/__init__.py index 355ed857592..24b4898865a 100644 --- a/examples/models/gemma4/__init__.py +++ b/examples/models/gemma4/__init__.py @@ -6,8 +6,9 @@ from executorch.examples.models.gemma4.text_decoder import ( convert_hf_to_custom, + convert_weights, Gemma4Config, Gemma4Model, ) -__all__ = ["Gemma4Config", "Gemma4Model", "convert_hf_to_custom"] +__all__ = ["Gemma4Config", "Gemma4Model", "convert_hf_to_custom", "convert_weights"] diff --git a/examples/models/gemma4/config/e2b_config.json b/examples/models/gemma4/config/e2b_config.json index 698f99131e6..d1401f9c229 100644 --- a/examples/models/gemma4/config/e2b_config.json +++ b/examples/models/gemma4/config/e2b_config.json @@ -22,6 +22,9 @@ "hidden_size_per_layer_input": 256, "num_kv_shared_layers": 20, "tie_word_embeddings": true, + "post_attention_norm": true, + "post_ffn_norm": true, + "use_ffn_norm": true, "layer_types": [ "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", diff --git a/examples/models/gemma4/text_decoder/__init__.py b/examples/models/gemma4/text_decoder/__init__.py index 4845a7e2c01..5ecde311ee8 100644 --- a/examples/models/gemma4/text_decoder/__init__.py +++ b/examples/models/gemma4/text_decoder/__init__.py @@ -5,7 +5,7 @@ # pyre-unsafe # LICENSE file in the root directory of this source tree. -from .convert_weights import convert_hf_to_custom # noqa: F401 +from .convert_weights import convert_hf_to_custom, convert_weights # noqa: F401 from .gemma4_attention import ( # noqa: F401 apply_rotary_emb, apply_rotary_emb_single, diff --git a/examples/models/gemma4/text_decoder/convert_weights.py b/examples/models/gemma4/text_decoder/convert_weights.py index e6ba1488a79..67d2f3de111 100644 --- a/examples/models/gemma4/text_decoder/convert_weights.py +++ b/examples/models/gemma4/text_decoder/convert_weights.py @@ -21,11 +21,41 @@ import torch +from executorch.examples.models.checkpoint import get_mapped_key + from .gemma4_config import Gemma4Config logger = logging.getLogger(__name__) +# Weight mappings from Gemma 4's checkpoint to ExecuTorch's transformer parameters. +_GEMMA4_TO_EXECUTORCH = { + "model.language_model.embed_tokens.weight": "tok_embeddings.weight", + "model.language_model.embed_tokens_per_layer.weight": "embed_tokens_per_layer.weight", + "model.language_model.per_layer_model_projection.weight": "per_layer_model_projection.weight", + "model.language_model.per_layer_projection_norm.weight": "per_layer_projection_norm.weight", + "model.language_model.norm.weight": "norm.weight", + "model.language_model.lm_head.weight": "output.weight", + "model.language_model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight", + "model.language_model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight", + "model.language_model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight", + "model.language_model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight", + "model.language_model.layers.{}.self_attn.q_norm.weight": "layers.{}.attention.q_norm_fn.weight", + "model.language_model.layers.{}.self_attn.k_norm.weight": "layers.{}.attention.k_norm_fn.weight", + "model.language_model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight", + "model.language_model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight", + "model.language_model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight", + "model.language_model.layers.{}.layer_scalar": "layers.{}.layer_scalar", + "model.language_model.layers.{}.per_layer_input_gate.weight": "layers.{}.per_layer_input_gate.weight", + "model.language_model.layers.{}.per_layer_projection.weight": "layers.{}.per_layer_projection.weight", + "model.language_model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight", + "model.language_model.layers.{}.post_attention_layernorm.weight": "layers.{}.post_attention_norm.weight", + "model.language_model.layers.{}.pre_feedforward_layernorm.weight": "layers.{}.ffn_norm.weight", + "model.language_model.layers.{}.post_feedforward_layernorm.weight": "layers.{}.post_ffn_norm.weight", + "model.language_model.layers.{}.post_per_layer_input_norm.weight": "layers.{}.post_per_layer_input_norm.weight", +} + + def _download_manifold_file(manifold_path: str, local_path: Path) -> None: """Download a file from Manifold to a local path.""" import io @@ -279,6 +309,43 @@ def convert_hf_to_custom( return converted_state_dict +def convert_weights(input_dir: str, output_file: str) -> None: + logger.info(f"Loading weights from {input_dir}") + + hf_state_dict = _load_safetensors_weights(input_dir) + + converted_state_dict = {} + unmapped_keys = [] + for hf_key, tensor in hf_state_dict.items(): + try: + qnn_key = get_mapped_key(hf_key, _GEMMA4_TO_EXECUTORCH) + except Exception: + unmapped_keys.append(hf_key) + logger.warning(f"Unmapped key: {hf_key}") + continue + converted_state_dict[qnn_key] = tensor.to(torch.float32) + logger.debug(f"Mapped {hf_key} -> {qnn_key}") + + lm_head_key = "output.weight" + embed_tokens_key = "tok_embeddings.weight" + if ( + lm_head_key not in converted_state_dict + and embed_tokens_key in converted_state_dict + ): + logger.info( + f"Using tied embeddings: copying {embed_tokens_key} to {lm_head_key}" + ) + converted_state_dict[lm_head_key] = converted_state_dict[ + embed_tokens_key + ].clone() + + logger.info(f"Converted {len(converted_state_dict)} weights") + if unmapped_keys: + logger.warning(f"Unmapped keys: {len(unmapped_keys)}") + + torch.save(converted_state_dict, output_file) + + def verify_conversion( hf_checkpoint_path: str, custom_state_dict: Dict[str, torch.Tensor], diff --git a/examples/models/llama/evaluate/eager_eval.py b/examples/models/llama/evaluate/eager_eval.py index 5c129e1c250..cefb951281c 100644 --- a/examples/models/llama/evaluate/eager_eval.py +++ b/examples/models/llama/evaluate/eager_eval.py @@ -50,6 +50,8 @@ def eot_token_id(self): @property def prefix_token_id(self): + if hasattr(self._tokenizer, "bos_id"): + return self._tokenizer.bos_id return self.eot_token_id @property diff --git a/examples/models/llama/model_args.py b/examples/models/llama/model_args.py index 7fca7523b82..86639d5f538 100644 --- a/examples/models/llama/model_args.py +++ b/examples/models/llama/model_args.py @@ -11,6 +11,7 @@ class ActFn(Enum): SILU = "silu" GELU = "gelu" GELU_APPROX = "gelu_approx" + GELU_APPROX_TANH = "gelu_pytorch_tanh" @classmethod def from_string(cls, value: str) -> "ActFn": @@ -29,7 +30,7 @@ def get_function(self): return F.silu elif self == ActFn.GELU: return F.gelu - elif self == ActFn.GELU_APPROX: + elif self == ActFn.GELU_APPROX or self == ActFn.GELU_APPROX_TANH: return partial(F.gelu, approximate="tanh") else: raise ValueError(f"Unsupported activation function: {self}") @@ -172,6 +173,13 @@ class ModelArgs: # gemma2 attn and output soft capping final_logit_softcapping: Optional[float] = None attn_logit_softcapping: Optional[float] = None + # gemma4: head_dim override for full-attention layers (sliding layers keep head_dim). + global_head_dim: Optional[int] = None + # gemma4: per layer embedding + hidden_size_per_layer_input: Optional[int] = None + vocab_size_per_layer_input: Optional[int] = None + # gemma4: cross-decoder (KV-shared) layers use 2x hidden_dim in FeedForward. + use_double_wide_mlp: bool = False # rlformers forward-pass features for on-device model parity normalize_tok_embeddings: bool = False @@ -225,3 +233,49 @@ def find_multiple(n: int, k: int) -> int: # Convert string act_fn to enum if needed if isinstance(self.act_fn, str): self.act_fn = ActFn.from_string(self.act_fn) + + def get_layer_type(self, layer_idx: int) -> str: + return self.layer_types[layer_idx] + + def is_sliding_attention(self, layer_idx: int) -> bool: + return ( + self.layer_types and self.get_layer_type(layer_idx) == "sliding_attention" + ) + + def get_head_dim(self, layer_idx: int) -> int: + if not self.global_head_dim: + return self.head_dim + if self.is_sliding_attention(layer_idx): + return self.head_dim + return self.global_head_dim + + def is_kv_shared_layer(self, layer_idx: int) -> bool: + first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers + return layer_idx >= first_kv_shared_layer_idx and first_kv_shared_layer_idx > 0 + + def get_intermediate_size(self, layer_idx: int) -> int: + if self.use_double_wide_mlp and self.is_kv_shared_layer(layer_idx): + return self.hidden_dim * 2 + return self.hidden_dim + + def get_kv_shared_layer_index(self, layer_idx: int) -> Optional[int]: + if not self.is_kv_shared_layer(layer_idx): + return None + first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers + current_layer_type = self.get_layer_type(layer_idx) + for i in range(first_kv_shared_layer_idx - 1, -1, -1): + if self.get_layer_type(i) == current_layer_type: + return i + return None + + def is_kv_donor_layer(self, layer_idx: int) -> bool: + if self.is_kv_shared_layer(layer_idx): + return False + first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers + if first_kv_shared_layer_idx <= 0: + return False + current_layer_type = self.get_layer_type(layer_idx) + for i in range(first_kv_shared_layer_idx - 1, -1, -1): + if self.get_layer_type(i) == current_layer_type: + return layer_idx == i + return False diff --git a/examples/qualcomm/oss_scripts/llama/README.md b/examples/qualcomm/oss_scripts/llama/README.md index 4522582f2f6..6892c489df8 100644 --- a/examples/qualcomm/oss_scripts/llama/README.md +++ b/examples/qualcomm/oss_scripts/llama/README.md @@ -14,6 +14,7 @@ This file provides you the instructions to run LLM Decoder model, VLM model, and 1. Gemma 2B 1. Gemma2 2B 1. Gemma3 1B + 1. Gemma4 E2B 1. GLM 1.5B 1. Granite3.3 2B 1. Phi4-mini-instruct @@ -157,6 +158,12 @@ Default example using hybrid mode python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma3-1b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1 ``` +#### Gemma4 E2B +Default example using hybrid mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma4-e2b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1 --embedding-quantize '4,32' +``` + #### GLM 1.5B Default example using hybrid mode ```bash diff --git a/examples/qualcomm/oss_scripts/llama/__init__.py b/examples/qualcomm/oss_scripts/llama/__init__.py index 2e0b2278909..6be8e15a380 100644 --- a/examples/qualcomm/oss_scripts/llama/__init__.py +++ b/examples/qualcomm/oss_scripts/llama/__init__.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import os from abc import ABC from dataclasses import dataclass @@ -17,6 +18,7 @@ from executorch.examples.models.gemma import convert_weights as convert_gemma_weights from executorch.examples.models.gemma2 import convert_weights as convert_gemma2_weights from executorch.examples.models.gemma3 import convert_weights as convert_gemma3_weights +from executorch.examples.models.gemma4 import convert_weights as convert_gemma4_weights from executorch.examples.models.glm import convert_weights as convert_glm_weights from executorch.examples.models.granite import ( @@ -59,9 +61,10 @@ SmolVLMEncoder, VisionModalityConfig, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( +from executorch.examples.qualcomm.oss_scripts.llama.model import ( LlamaModel, LlamaModelWithoutEmbedding, + ModelArgs, MultiScopeAwareLlamaModel, ) @@ -69,6 +72,7 @@ CodegenQuantRecipe, Gemma2QuantRecipe, Gemma3QuantRecipe, + Gemma4QuantRecipe, Gemma_2BQuantRecipe, GLM_1_5B_InstructQuantRecipe, Granite_3_3_2B_InstructQuantRecipe, @@ -100,6 +104,7 @@ "smolvlm_500m_instruct": LlamaModelWithoutEmbedding, "internvl3_1b": LlamaModelWithoutEmbedding, "gemma2-2b": MultiScopeAwareLlamaModel, + "gemma4-e2b": MultiScopeAwareLlamaModel, } @@ -602,3 +607,56 @@ class SmolVLM_500M(LLMModelConfig): r2 = False r3 = False quant_recipe = SmolVLMQuantRecipe + + +def load_gemma4_model_args(params_path: str) -> ModelArgs: + with open(params_path) as f: + config_dict = json.load(f) + text_config = config_dict.get("text_config", config_dict) + + return ModelArgs( + dim=text_config["hidden_size"], + n_layers=text_config["num_hidden_layers"], + n_heads=text_config["num_attention_heads"], + n_kv_heads=text_config["num_key_value_heads"], + head_dim=text_config["head_dim"], + global_head_dim=text_config["global_head_dim"], + vocab_size=text_config["vocab_size"], + hidden_dim=text_config["intermediate_size"], + norm_eps=text_config["rms_norm_eps"], + max_position_embeddings=text_config["max_position_embeddings"], + rope_theta=text_config["rope_theta"], + local_rope_theta=text_config["rope_local_base_freq"], + partial_rotary_factor=text_config["partial_rotary_factor"], + sliding_window=text_config["sliding_window"], + layer_types=text_config["layer_types"], + final_logit_softcapping=text_config["final_logit_softcapping"], + hidden_size_per_layer_input=text_config["hidden_size_per_layer_input"], + vocab_size_per_layer_input=text_config["vocab_size_per_layer_input"], + num_kv_shared_layers=text_config["num_kv_shared_layers"], + use_double_wide_mlp=text_config["use_double_wide_mlp"], + post_attention_norm=text_config["post_attention_norm"], + post_ffn_norm=text_config["post_ffn_norm"], + use_ffn_norm=text_config["use_ffn_norm"], + act_fn=text_config["hidden_activation"], + model_architecture=config_dict["model_type"], + ) + + +@register_llm_model("gemma4-e2b") +@dataclass(init=False, frozen=True) +class Gemma4_E2B(LLMModelConfig): + repo_id: str = "google/gemma-4-e2b-it" + params_path: str = os.path.join( + BASE_DIR, "../../../models/gemma4/config/e2b_config.json" + ) + convert_weights = convert_gemma4_weights + transform_weight = False + instruct_model = True + num_sharding = 4 + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + quant_recipe = Gemma4QuantRecipe diff --git a/examples/qualcomm/oss_scripts/llama/decoder_constants.py b/examples/qualcomm/oss_scripts/llama/decoder_constants.py index cccf368459b..26e5aee1ad7 100644 --- a/examples/qualcomm/oss_scripts/llama/decoder_constants.py +++ b/examples/qualcomm/oss_scripts/llama/decoder_constants.py @@ -63,4 +63,5 @@ "glm-1_5b": "glm", "smolvlm_500m_instruct": "smolvlm", "internvl3_1b": "internvl3", + "gemma4-e2b": "gemma4", } diff --git a/examples/qualcomm/oss_scripts/llama/encoder/encoder_config.py b/examples/qualcomm/oss_scripts/llama/encoder/encoder_config.py index b53f4bda689..46f64b838a3 100644 --- a/examples/qualcomm/oss_scripts/llama/encoder/encoder_config.py +++ b/examples/qualcomm/oss_scripts/llama/encoder/encoder_config.py @@ -13,10 +13,10 @@ InternVL3EncoderQuantRecipe, SmolVLMEncoderQuantRecipe, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.audio_encoder import ( +from executorch.examples.qualcomm.oss_scripts.llama.model.encoders.audio import ( GraniteSpeechCTCEncoderWrapper, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.vision_encoder import ( +from executorch.examples.qualcomm.oss_scripts.llama.model.encoders.vision import ( Idefics3VisionEncoder, InternVL3VisionEncoder, ) diff --git a/examples/qualcomm/oss_scripts/llama/llama.py b/examples/qualcomm/oss_scripts/llama/llama.py index edf719a6fe5..ad4e5257fe5 100755 --- a/examples/qualcomm/oss_scripts/llama/llama.py +++ b/examples/qualcomm/oss_scripts/llama/llama.py @@ -666,6 +666,11 @@ def export_llama(args) -> None: "Pass --calib_tasks to match the previous --tasks behavior." ) raise RuntimeError("Please provide --eval_tasks to eval perplexity") + if args.decoder_model == "gemma4-e2b" and not args.embedding_quantize: + raise RuntimeError( + "gemma4-e2b requires --embedding-quantize: " + "per-layer embedding is too large for HTP quantization." + ) assert ( args.decoder_model in SUPPORTED_LLM_MODELS ), f"Unknown decoder_model: {args.decoder_model}." diff --git a/examples/qualcomm/oss_scripts/llama/model/__init__.py b/examples/qualcomm/oss_scripts/llama/model/__init__.py index 8300afd4df3..c9e3706a20a 100644 --- a/examples/qualcomm/oss_scripts/llama/model/__init__.py +++ b/examples/qualcomm/oss_scripts/llama/model/__init__.py @@ -4,13 +4,68 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from .apply_rope import ROTARY_EMB_REGISTRY -from .feed_forward import FeedForward_REGISTRY -from .layernorm import NORM_REGISTRY +from executorch.examples.models.llama.model_args import ModelArgs + +from .text_decoder import ( + ATTENTION_REGISTRY, + AttentionSinkRope, + DECODER_LAYER_REGISTRY, + FEED_FORWARD_REGISTRY, + FeedForward, + FeedForwardBase, + Gemma4Attention, + LayerNorm, + LlamaAttention, + LlamaDecoderLayer, + LlamaModel, + LlamaModelWithoutEmbedding, + MultiScopeAwareLlamaModel, + Norm, + NORM_REGISTRY, + register_attention, + register_decoder_layer, + register_feed_forward, + register_norm, + register_rope, + register_rope_precompute, + repeat_kv, + RMSNorm, + ROPE_PRECOMPUTE_REGISTRY, + ROPE_REGISTRY, + RopeFreqs, + TokenEmbedding, + YOCOCrossDecoderLayer, +) __all__ = [ - "FeedForward_REGISTRY", - "ROTARY_EMB_REGISTRY", + "ModelArgs", "NORM_REGISTRY", + "Norm", + "LayerNorm", + "RMSNorm", + "register_norm", + "ROPE_REGISTRY", + "ROPE_PRECOMPUTE_REGISTRY", + "RopeFreqs", + "register_rope", + "register_rope_precompute", + "FEED_FORWARD_REGISTRY", + "FeedForward", + "FeedForwardBase", + "register_feed_forward", + "repeat_kv", + "LlamaAttention", + "LlamaDecoderLayer", + "TokenEmbedding", + "AttentionSinkRope", + "LlamaModel", + "LlamaModelWithoutEmbedding", + "MultiScopeAwareLlamaModel", + "Gemma4Attention", + "ATTENTION_REGISTRY", + "register_attention", + "YOCOCrossDecoderLayer", + "DECODER_LAYER_REGISTRY", + "register_decoder_layer", ] diff --git a/examples/qualcomm/oss_scripts/llama/model/apply_rope.py b/examples/qualcomm/oss_scripts/llama/model/apply_rope.py deleted file mode 100644 index 9ad23eb5c41..00000000000 --- a/examples/qualcomm/oss_scripts/llama/model/apply_rope.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. -from typing import Callable, Dict - -import torch - - -ROTARY_EMB_REGISTRY: Dict[str, Callable] = {} - - -def register_rotary_emb(name: str): - """Register a rotary embedding function.""" - - def decorator(fn: Callable): - ROTARY_EMB_REGISTRY[name] = fn - return fn - - return decorator - - -@register_rotary_emb("partial") -def apply_partial_rotary_emb_single(x, freqs_cos, freqs_sin): - if x.dim() == 4: - freqs_cos = freqs_cos[None, None, :, :] - freqs_sin = freqs_sin[None, None, :, :] - rotary_dim = freqs_cos.shape[-1] * 2 - x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] - x_r, x_i = x_rot[..., : x_rot.shape[-1] // 2], x_rot[..., x_rot.shape[-1] // 2 :] - x_out_r = x_r * freqs_cos - x_i * freqs_sin - x_out_i = x_r * freqs_sin + x_i * freqs_cos - x_rotated = torch.cat([x_out_r, x_out_i], dim=-1) - return torch.cat([x_rotated, x_pass], dim=-1) - - -@register_rotary_emb("default") -def apply_rotary_emb_single(x, freqs_cos, freqs_sin): - # The implementation of RoPE in HuggingFace processes query and key with two half instead of interleaved way. - # The main difference is stride in StrideSlice op. For interleaved way, stride is two which is not friendly for HTP backend. - # Ref: https://github.com/huggingface/transformers/issues/25199 - x_r, x_i = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] - # broadcast for batch_prefill mode input x - if x.dim() == 4: - freqs_cos = freqs_cos[None, None, :, :] - freqs_sin = freqs_sin[None, None, :, :] - x_out_r = x_r * freqs_cos - x_i * freqs_sin - x_out_i = x_r * freqs_sin + x_i * freqs_cos - - x_out = torch.cat([x_out_r, x_out_i], dim=-1) - return x_out diff --git a/examples/qualcomm/oss_scripts/llama/model/encoders/__init__.py b/examples/qualcomm/oss_scripts/llama/model/encoders/__init__.py new file mode 100644 index 00000000000..32546682aaf --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/encoders/__init__.py @@ -0,0 +1,15 @@ +# 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. + +from .audio import GraniteSpeechCTCEncoderWrapper +from .vision import Idefics3VisionEncoder, InternVL3VisionEncoder + + +__all__ = [ + "GraniteSpeechCTCEncoderWrapper", + "Idefics3VisionEncoder", + "InternVL3VisionEncoder", +] diff --git a/examples/qualcomm/oss_scripts/llama/model/audio_encoder.py b/examples/qualcomm/oss_scripts/llama/model/encoders/audio.py similarity index 100% rename from examples/qualcomm/oss_scripts/llama/model/audio_encoder.py rename to examples/qualcomm/oss_scripts/llama/model/encoders/audio.py diff --git a/examples/qualcomm/oss_scripts/llama/model/vision_encoder.py b/examples/qualcomm/oss_scripts/llama/model/encoders/vision.py similarity index 100% rename from examples/qualcomm/oss_scripts/llama/model/vision_encoder.py rename to examples/qualcomm/oss_scripts/llama/model/encoders/vision.py diff --git a/examples/qualcomm/oss_scripts/llama/model/static_llama.py b/examples/qualcomm/oss_scripts/llama/model/static_llama.py deleted file mode 100755 index d57ada7fef6..00000000000 --- a/examples/qualcomm/oss_scripts/llama/model/static_llama.py +++ /dev/null @@ -1,1088 +0,0 @@ -# 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. - -# TODO: reenable pyre after fixing the issues -# pyre-ignore-all-errors - -import math -from typing import List, Optional, Tuple - -import scipy -import torch -import torch.nn as nn - -from executorch.examples.models.llama.model_args import ModelArgs -from executorch.examples.models.llama.rope import ( - hf_precompute_freqs_cis, - precompute_freqs_cis, -) -from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import ( - AttentionMask, - CausalAttentionMask, - SlidingWindowAttentionMask, -) -from executorch.examples.qualcomm.oss_scripts.llama.model import ( - FeedForward_REGISTRY, - NORM_REGISTRY, - ROTARY_EMB_REGISTRY, -) - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand( - batch, num_key_value_heads, n_rep, slen, head_dim - ) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -class AttentionSinkRope(nn.Module): - def __init__( - self, - config: ModelArgs, - sink_size: int, - eviction_batch_size: int, - ar_len: int, - **kwargs, - ): - super().__init__() - self.config = config - self.sink_size = sink_size - self.eviction_batch_size = eviction_batch_size - self.n_layers = config.n_layers - self.original_position = eviction_batch_size + sink_size - self.new_position = sink_size - self.num_to_keep = ( - config.max_context_len - sink_size - eviction_batch_size - ar_len - ) - self.evict_k_cache_shape = [ - config.max_batch_size, - config.n_kv_heads, - config.head_dim, - eviction_batch_size, - ] - self.evict_v_cache_shape = [ - config.max_batch_size, - config.n_kv_heads, - eviction_batch_size, - config.head_dim, - ] - self.kv_cache_shape = { - # single head, k input - "k": ( - config.max_batch_size, - config.n_kv_heads, - config.head_dim, - config.max_context_len - ar_len, - ), - # single head, v input - "v": ( - config.max_batch_size, - config.n_kv_heads, - config.max_context_len - ar_len, - config.head_dim, - ), - } - - if getattr(config, "enable_r3", False): - self.register_buffer( - "r3_weight", - torch.tensor( - scipy.linalg.hadamard(config.head_dim, dtype=float) - / math.sqrt(config.head_dim), - dtype=torch.float32, - device="cpu", - ), - persistent=False, - ) - - if config.partial_rotary_factor < 1: - self.apply_rope_emb = ROTARY_EMB_REGISTRY["partial"] - else: - self.apply_rope_emb = ROTARY_EMB_REGISTRY["default"] - - if config.use_hf_rope: - freqs_cos, freqs_sin = hf_precompute_freqs_cis( - config.head_dim, - config.max_context_len, - config.rope_freq_base, - config.partial_rotary_factor, - ) - freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] - freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] - else: - freqs_cos, freqs_sin = precompute_freqs_cis( - config.head_dim, - config.max_context_len, - config.rope_freq_base, - config.use_scaled_rope, - config.rope_scale_factor, - ) - original_freqs_cos = freqs_cos.narrow( - 0, self.original_position, self.num_to_keep - ) - original_freqs_sin = freqs_sin.narrow( - 0, self.original_position, self.num_to_keep - ) - new_freqs_cos = freqs_cos.narrow(0, self.new_position, self.num_to_keep) - new_freqs_sin = freqs_sin.narrow(0, self.new_position, self.num_to_keep) - rerotation_cos = ( - new_freqs_cos * original_freqs_cos + new_freqs_sin * original_freqs_sin - ) - rerotation_sin = ( - new_freqs_sin * original_freqs_cos - new_freqs_cos * original_freqs_sin - ) - self.register_buffer("rerotation_cos", rerotation_cos, persistent=False) - self.register_buffer("rerotation_sin", rerotation_sin, persistent=False) - - self.sliding_window = kwargs.get("sliding_window", False) - if self.sliding_window: - # Get attention type for each layer - self.layer_types = kwargs["layer_types"] - # Get local freq base for sliding attention - rope_freq_base = kwargs["rope_local_base_freq"] - local_freqs_cos, local_freqs_sin = hf_precompute_freqs_cis( - config.head_dim, - config.max_context_len, - rope_freq_base, - config.partial_rotary_factor, - ) - local_freqs_cos = local_freqs_cos[:, : local_freqs_cos.shape[-1] // 2] - local_freqs_sin = local_freqs_sin[:, : local_freqs_sin.shape[-1] // 2] - local_original_freqs_cos = local_freqs_cos.narrow( - 0, self.original_position, self.num_to_keep - ) - local_original_freqs_sin = local_freqs_sin.narrow( - 0, self.original_position, self.num_to_keep - ) - local_new_freqs_cos = local_freqs_cos.narrow( - 0, self.new_position, self.num_to_keep - ) - local_new_freqs_sin = local_freqs_sin.narrow( - 0, self.new_position, self.num_to_keep - ) - local_rerotation_cos = ( - local_new_freqs_cos * local_original_freqs_cos - + local_new_freqs_sin * local_original_freqs_sin - ) - local_rerotation_sin = ( - local_new_freqs_sin * local_original_freqs_cos - - local_new_freqs_cos * local_original_freqs_sin - ) - self.register_buffer( - "local_rerotation_cos", local_rerotation_cos, persistent=False - ) - self.register_buffer( - "local_rerotation_sin", local_rerotation_sin, persistent=False - ) - - def forward(self, k_caches: List[torch.Tensor], v_caches: List[torch.Tensor]): - """ - Rerotate k_cache from original_position to new_position, and return the kv cache after eviction. This is done by rerotating - k_cache with (new_position * theta - original_position * theta) with the following matrix: - (cos(delta), -sin(delta) - sin(delta), cos(delta)) - where delta = new_position * theta - original_position * theta - - Based on https://github.com/huggingface/transformers/pull/26681 - """ - - output_k_caches, output_v_caches = [], [] - for ind, (k_cache, v_cache) in enumerate(zip(k_caches, v_caches)): - # k_cache: (batch_size, n_kv_heads, head_dim, seq_len) - # v_cache: (batch_size, n_kv_heads, seq_len, head_dim) - k_dim_to_slice = 3 - v_dim_to_slice = 2 - - k_to_keep = k_cache.narrow( - k_dim_to_slice, - self.original_position, - self.num_to_keep, - ) - k_to_keep = k_to_keep.transpose(2, 3) - if getattr(self.config, "enable_r3", False): - # We need to revert the key from spin quant before applying RoPE - k_to_keep = torch.matmul(k_to_keep, self.r3_weight.T) - - if self.sliding_window and self.layer_types[ind] == "sliding_attention": - k_to_keep = self.apply_rope_emb( - k_to_keep, self.local_rerotation_cos, self.local_rerotation_sin - ) - else: - k_to_keep = self.apply_rope_emb( - k_to_keep, self.rerotation_cos, self.rerotation_sin - ) - if getattr(self.config, "enable_r3", False): - k_to_keep = torch.matmul(k_to_keep, self.r3_weight) - k_to_keep = k_to_keep.transpose(2, 3) - new_k_cache = torch.cat( - [ - k_cache.narrow(k_dim_to_slice, 0, self.sink_size), - k_to_keep, - torch.zeros(self.evict_k_cache_shape), - ], - dim=k_dim_to_slice, - ) - - new_v_cache = torch.cat( - [ - v_cache.narrow(v_dim_to_slice, 0, self.sink_size), - v_cache.narrow( - v_dim_to_slice, - self.original_position, - self.num_to_keep, - ), - torch.zeros(self.evict_v_cache_shape), - ], - dim=v_dim_to_slice, - ) - - output_k_caches.append(new_k_cache) - output_v_caches.append(new_v_cache) - - return output_k_caches, output_v_caches - - def get_example_inputs(self): - k_cache, v_cache = [], [] - - for _ in range(self.n_layers): - k_cache.append(torch.zeros(self.kv_cache_shape["k"])) - v_cache.append(torch.zeros(self.kv_cache_shape["v"])) - return k_cache, v_cache - - def get_metadata(self): - return { - "get_eviction_batch_size": self.eviction_batch_size, - "get_max_context_len": self.config.max_context_len, - "get_sink_size": self.sink_size, - } - - -class LlamaAttention(nn.Module): - def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): - super().__init__() - self.config = config - self.dim = config.dim - self.n_heads = config.n_heads - self.head_dim = config.head_dim - self.n_kv_heads = config.n_kv_heads - self.num_key_value_groups = config.n_heads // self.n_kv_heads - self.max_seq_len = config.max_seq_len - self.max_context_len = config.max_context_len - self.use_kv_cache = config.use_kv_cache - self.output_new_cache_only = output_new_cache_only - self.enable_masked_softmax = getattr(config, "enable_masked_softmax", False) - self.use_qk_norm = config.use_qk_norm - self.qk_norm_before_rope = config.qk_norm_before_rope - # If None, assume each layer uses rope - self.use_rope = ( - config.no_rope_layer_interval is None - or (layer_idx + 1) % config.no_rope_layer_interval - ) - - if self.use_qk_norm: - q_norm_dim = self.head_dim - k_norm_dim = self.head_dim - self.q_norm_fn = torch.nn.RMSNorm(q_norm_dim, eps=config.norm_eps) - self.k_norm_fn = torch.nn.RMSNorm(k_norm_dim, eps=config.norm_eps) - - if config.partial_rotary_factor < 1: - self.apply_rope_emb = ROTARY_EMB_REGISTRY["partial"] - else: - self.apply_rope_emb = ROTARY_EMB_REGISTRY["default"] - - self.wq = nn.Linear( - self.dim, - self.n_heads * self.head_dim, - bias=getattr(config, "attention_qkv_bias", False), - ) - self.wk = nn.Linear( - self.dim, - self.n_kv_heads * self.head_dim, - bias=getattr(config, "attention_qkv_bias", False), - ) - self.wv = nn.Linear( - self.dim, - self.n_kv_heads * self.head_dim, - bias=getattr(config, "attention_qkv_bias", False), - ) - self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) - - self.attn_softmax = torch.nn.Softmax(dim=-1) - - self.scale = ( - float(self.head_dim) ** 0.5 - if config.attention_multiplier is None - else 1.0 / config.attention_multiplier - ) - - # gemma 2 uses soft-capping on attention and logits - self.attn_logit_softcapping = config.attn_logit_softcapping - - if getattr(config, "enable_r3", False): - self.register_buffer( - "r3_weight", - torch.tensor( - scipy.linalg.hadamard(self.head_dim, dtype=float) - / math.sqrt(self.head_dim), - dtype=torch.float32, - device="cpu", - ), - persistent=False, - ) - - def prepare_attention_conv(self): - self.wq_conv = nn.Conv2d( - self.dim, - self.n_heads * self.head_dim, - 1, - bias=getattr(self.config, "attention_qkv_bias", False), - ) - self.wk_conv = nn.Conv2d( - self.dim, - self.n_kv_heads * self.head_dim, - 1, - bias=getattr(self.config, "attention_qkv_bias", False), - ) - self.wv_conv = nn.Conv2d( - self.dim, - self.n_kv_heads * self.head_dim, - 1, - bias=getattr(self.config, "attention_qkv_bias", False), - ) - self.wo_conv = nn.Conv2d(self.n_heads * self.head_dim, self.dim, 1, bias=False) - - self.forward_no_conv = self.forward - self.forward = self.forward_attention_conv - - self.wq_conv.weight.data.copy_(self.wq.weight[:, :, None, None]) - if self.wq_conv.bias is not None: - self.wq_conv.bias.data.copy_(self.wq.bias) - self.wk_conv.weight.data.copy_(self.wk.weight[:, :, None, None]) - if self.wk_conv.bias is not None: - self.wk_conv.bias.data.copy_(self.wk.bias) - self.wv_conv.weight.data.copy_(self.wv.weight[:, :, None, None]) - if self.wv_conv.bias is not None: - self.wv_conv.bias.data.copy_(self.wv.bias) - self.wo_conv.weight.data.copy_(self.wo.weight[:, :, None, None]) - - del self.wq - del self.wk - del self.wv - del self.wo - - def forward_attention_conv( - self, - hidden_states: torch.Tensor, - freqs_cos: torch.Tensor, - freqs_sin: torch.Tensor, - atten_mask: torch.Tensor, - k_caches: List[torch.Tensor], - v_caches: List[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - bsz, seq_len, _ = hidden_states.shape - hidden_states = torch.reshape( - hidden_states, (bsz, seq_len, 1, self.dim) - ).transpose(1, 3) - - q = self.wq_conv(hidden_states) - k = self.wk_conv(hidden_states) - v = self.wv_conv(hidden_states) - q = q.permute(0, 3, 1, 2).squeeze(-1) - k = k.permute(0, 3, 1, 2).squeeze(-1) - v = v.permute(0, 3, 1, 2).squeeze(-1) - q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) - k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) - v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) - - if self.use_qk_norm and self.qk_norm_before_rope: - q = self.q_norm_fn(q) - k = self.k_norm_fn(k) - - if self.use_rope: - q = self.apply_rope_emb(q, freqs_cos, freqs_sin) - k = self.apply_rope_emb(k, freqs_cos, freqs_sin) - - if self.use_qk_norm and not self.qk_norm_before_rope: - q = self.q_norm_fn(q) - k = self.k_norm_fn(k) - if getattr(self.config, "enable_r3", False): - q = torch.matmul(q, self.r3_weight) - k = torch.matmul(k, self.r3_weight) - k = k.transpose(2, 3) - - kh, vh = None, None - # kv cache mode - if self.use_kv_cache: - kh = torch.cat([k_caches, k], dim=-1) - vh = torch.cat([v_caches, v], dim=2) - # batch_prefill mode - else: - kh = k - vh = v - - kh = repeat_kv(kh, self.num_key_value_groups) - vh = repeat_kv(vh, self.num_key_value_groups) - - attn = q @ kh - attn = attn / self.scale - if self.enable_masked_softmax: - attn_min = torch.amin(attn, dim=-1, keepdim=True) - minus_value = -20 - attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) - else: - attn = attn + atten_mask - attn = self.attn_softmax(attn) - y = attn @ vh - y = y.transpose(1, 2) - y = y.reshape(bsz, seq_len, 1, -1).transpose(1, 3) - y = self.wo_conv(y) - y = y.transpose(1, 3) - y = y.reshape(bsz, seq_len, -1) - - if self.output_new_cache_only: - return y, [k], [v] - - return y, [kh], [vh] - - def forward( - self, - hidden_states: torch.Tensor, - freqs_cos: torch.Tensor, - freqs_sin: torch.Tensor, - atten_mask: torch.Tensor, - k_caches: List[torch.Tensor], - v_caches: List[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - bsz, seq_len, _ = hidden_states.shape - - q = self.wq(hidden_states) - k = self.wk(hidden_states) - v = self.wv(hidden_states) - q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) - k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) - v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) - - if self.use_qk_norm and self.qk_norm_before_rope: - q = self.q_norm_fn(q) - k = self.k_norm_fn(k) - - if self.use_rope: - q = self.apply_rope_emb(q, freqs_cos, freqs_sin) - k = self.apply_rope_emb(k, freqs_cos, freqs_sin) - - if self.use_qk_norm and not self.qk_norm_before_rope: - q = self.q_norm_fn(q) - k = self.k_norm_fn(k) - if getattr(self.config, "enable_r3", False): - q = torch.matmul(q, self.r3_weight) - k = torch.matmul(k, self.r3_weight) - k = k.transpose(2, 3) - - kh, vh = None, None - # kv cache mode - if self.use_kv_cache: - kh = torch.cat([k_caches, k], dim=-1) - vh = torch.cat([v_caches, v], dim=2) - # batch_prefill mode - else: - kh = k - vh = v - - kh = repeat_kv(kh, self.num_key_value_groups) - vh = repeat_kv(vh, self.num_key_value_groups) - - attn = q @ kh - # gemma2-2b - if self.attn_logit_softcapping is not None: - attn = attn / self.attn_logit_softcapping - attn = torch.tanh(attn) - attn = attn * self.attn_logit_softcapping - attn = attn / self.scale - if self.enable_masked_softmax: - attn_min = torch.amin(attn, dim=-1, keepdim=True) - minus_value = -20 - attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) - else: - attn = attn + atten_mask - attn = self.attn_softmax(attn) - y = attn @ vh - y = y.transpose(1, 2).reshape(bsz, seq_len, -1) - y = self.wo(y) - - if self.output_new_cache_only: - return y, [k], [v] - - return y, [kh], [vh] - - -class FeedForward(nn.Module): - def __init__(self, args: ModelArgs): - super().__init__() - assert args.hidden_dim is not None - self.hidden_dim: int = args.hidden_dim - self.dim: int = args.dim - self.w1 = nn.Linear(self.dim, self.hidden_dim, bias=False) - self.w2 = nn.Linear(self.hidden_dim, self.dim, bias=False) - self.w3 = nn.Linear(self.dim, self.hidden_dim, bias=False) - self.act_fn = args.act_fn.get_function() - - def prepare_feedfoward_conv(self): - self.w1_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False) - self.w2_conv = nn.Conv2d(self.hidden_dim, self.dim, 1, bias=False) - self.w3_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False) - - self.forward_no_conv = self.forward - self.forward = self.forward_feedfoward_conv - self.w1_conv.weight.data.copy_(self.w1.weight[:, :, None, None]) - self.w2_conv.weight.data.copy_(self.w2.weight[:, :, None, None]) - self.w3_conv.weight.data.copy_(self.w3.weight[:, :, None, None]) - - del self.w1 - del self.w2 - del self.w3 - - def forward_feedfoward_conv(self, x): - bsz, _, _ = x.size() - x = torch.reshape(x, (bsz, -1, 1, self.dim)) - x = x.transpose(1, 3) # Transpose right before and after Conv - x = self.w2_conv(self.act_fn(self.w1_conv(x)) * self.w3_conv(x)) - x = x.transpose(1, 3) - x = torch.reshape(x, (bsz, -1, self.dim)) - return x - - def forward(self, x): - return self.w2(self.act_fn(self.w1(x)) * self.w3(x)) - - -class LlamaDecoderLayer(nn.Module): - def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): - super().__init__() - self.dim = config.dim - self.attention = LlamaAttention( - layer_idx=layer_idx, - config=config, - output_new_cache_only=output_new_cache_only, - ) - - self.feed_forward = FeedForward_REGISTRY.get( - config.model_architecture, FeedForward - )(config) - self.attention_norm = NORM_REGISTRY[config.norm_type]( - config.dim, eps=config.norm_eps - ) - self.ffn_norm = ( - NORM_REGISTRY[config.norm_type](config.dim, eps=config.norm_eps) - if config.use_ffn_norm - else None - ) - - self.post_attention_norm = ( - torch.nn.RMSNorm(config.dim, eps=config.norm_eps) - if config.post_attention_norm - else None - ) - self.residual_multiplier = config.residual_multiplier - self.post_ffn_norm = ( - torch.nn.RMSNorm(config.dim, eps=config.norm_eps) - if config.post_ffn_norm - else None - ) - - def forward( - self, - x: torch.Tensor, - freqs_cos: torch.Tensor, - freqs_sin: torch.Tensor, - atten_mask: torch.Tensor, - k_caches: List[torch.Tensor], - v_caches: List[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - - hidden_states = self.attention_norm(x) - h, k_cache, v_cache = self.attention( - hidden_states=hidden_states, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - atten_mask=atten_mask, - k_caches=k_caches, - v_caches=v_caches, - ) - if self.post_attention_norm: - h = self.post_attention_norm(h) - h = ( - x + h * self.residual_multiplier - if self.residual_multiplier is not None - else x + h - ) - hidden_states = hidden_states if self.ffn_norm is None else self.ffn_norm(h) - out = self.feed_forward(hidden_states) - if self.post_ffn_norm: - out = self.post_ffn_norm(out) - output = ( - h + out * self.residual_multiplier - if self.residual_multiplier is not None - else h + out - ) - - return output, k_cache, v_cache - - -class LlamaModel(nn.Module): - def __init__( - self, - config: ModelArgs, - ar_len=1, - output_new_cache_only=True, - output_cache=True, - use_i64_token=False, - **kwargs, - ): - super().__init__() - self.dim = config.dim - self.head_dim = config.head_dim - self.max_batch_size = config.max_batch_size - self.max_seq_len = config.max_seq_len - self.max_context_len = config.max_context_len - self.n_heads = config.n_heads - self.n_kv_heads = config.n_kv_heads - self.n_layers = config.n_layers - self.vocab_size = config.vocab_size - self.rope_freq_base = config.rope_freq_base - self.use_kv_cache = config.use_kv_cache - self.embedding_scale_factor = config.embedding_scale_factor - self.ar_len = ar_len - self.output_new_cache_only = output_new_cache_only - self.use_i64_token = use_i64_token - self.output_cache = output_cache - self.kv_io_bit_width = config.kv_io_bit_width - self.logits_scaling = config.logits_scaling - - self.layers = nn.ModuleList( - [ - LlamaDecoderLayer(i, config, self.output_new_cache_only) - for i in range(config.n_layers) - ] - ) - self.norm = NORM_REGISTRY[config.norm_type](config.dim, eps=config.norm_eps) - self.output = nn.Linear(config.dim, config.vocab_size, bias=config.output_bias) - - self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim) - if config.use_hf_rope: - freqs_cos, freqs_sin = hf_precompute_freqs_cis( - config.head_dim, - config.max_context_len, - config.rope_freq_base, - config.partial_rotary_factor, - ) - freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] - freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] - else: - freqs_cos, freqs_sin = precompute_freqs_cis( - config.head_dim, - config.max_context_len, - config.rope_freq_base, - config.use_scaled_rope, - config.rope_scale_factor, - ) - self.register_buffer("freqs_cos", freqs_cos, persistent=False) - self.register_buffer("freqs_sin", freqs_sin, persistent=False) - - def prepare_output_conv(self): - def forward_output_conv(x): - bsz, _, _ = x.size() - x = torch.reshape(x, (bsz, -1, 1, self.dim)) - x = x.transpose(1, 3) # Transpose right before and after Conv - x = self.output_conv(x) - x = x.transpose(1, 3) - x = torch.reshape(x, (bsz, -1, self.vocab_size)) - return x - - self.output_conv = nn.Conv2d(self.dim, self.vocab_size, 1, bias=False) - self.output_conv.weight.data.copy_(self.output.weight[:, :, None, None]) - - del self.output - self.output = forward_output_conv - - def forward( - self, - tokens: torch.Tensor, - atten_mask: torch.Tensor, - input_pos: Optional[torch.Tensor] = None, - *args, - ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: - output_k_cache = [] - output_v_cache = [] - # following tensors should be invariant across batches - freqs_cos = ( - self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos - ) - freqs_sin = ( - self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin - ) - - hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) - - for ind, decoder_layer in enumerate(self.layers): - k_caches = None - v_caches = None - if self.use_kv_cache: - offset_k = ind - offset_v = self.n_layers + offset_k - k_caches = args[offset_k] - v_caches = args[offset_v] - - hidden_states, k, v = decoder_layer( - hidden_states, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - atten_mask=atten_mask, - k_caches=k_caches, - v_caches=v_caches, - ) - output_k_cache.extend(k) - output_v_cache.extend(v) - - hidden_states = self.norm(hidden_states) - logits = self.output(hidden_states) - - if self.logits_scaling: - logits = logits / self.logits_scaling - - if self.output_cache: - return logits, output_k_cache, output_v_cache - return logits - - def get_example_inputs(self): - dtype = torch.int64 if self.use_i64_token else torch.int32 - tokens = torch.randint( - self.vocab_size, (self.max_batch_size, self.ar_len), dtype=dtype - ) - atten_mask = AttentionMask( - CausalAttentionMask(self.max_batch_size, self.ar_len, self.max_context_len) - ) - if self.use_kv_cache: - pos_ids = torch.zeros((self.max_batch_size, self.ar_len), dtype=torch.int32) - k_cache, v_cache = [], [] - for _ in range(self.n_layers): - # transpose first to decrease the runtime efforts - k_cache.append( - torch.zeros( - self.max_batch_size, - self.n_kv_heads, - self.head_dim, - self.max_context_len - self.ar_len, - ) - ) - v_cache.append( - torch.zeros( - self.max_batch_size, - self.n_kv_heads, - self.max_context_len - self.ar_len, - self.head_dim, - ) - ) - return ( - tokens, - atten_mask, - pos_ids, - k_cache, - v_cache, - ) - - return ( - tokens, - atten_mask, - ) - - def get_metadata(self): - return { - "get_ar_len": self.ar_len, - "get_bos_id": 1, - "get_eos_id": 2, - "get_dim": self.dim, - "get_head_dim": self.head_dim, - "get_max_batch_size": self.max_batch_size, - "get_max_seq_len": self.max_seq_len, - "get_max_context_len": self.max_context_len, - "get_n_bos": 1, - "get_n_eos": 1, - "get_n_kv_heads": self.n_kv_heads, - "get_n_layers": self.n_layers, - "get_vocab_size": self.vocab_size, - "get_use_kv_cache": self.use_kv_cache, - "get_kv_io_bit_width": self.kv_io_bit_width, - } - - -class LlamaModelWithoutEmbedding(LlamaModel): - def __init__( - self, - config: ModelArgs, - ar_len=1, - output_new_cache_only=True, - output_cache=True, - use_i64_token=False, - **kwargs, - ): - super().__init__( - config=config, - ar_len=ar_len, - output_new_cache_only=output_new_cache_only, - output_cache=output_cache, - use_i64_token=use_i64_token, - **kwargs, - ) - - # Set the audio/image token ID from keyword arguments. It defaults to None if not provided. - # If an ID is provided, it will be stored in the model's metadata. - self.audio_token_id = kwargs.get("audio_token_id", None) - self.image_token_id = kwargs.get("image_token_id", None) - - def forward( - self, - hidden_states: torch.Tensor, - atten_mask: torch.Tensor, - input_pos: Optional[torch.Tensor] = None, - *args, - ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: - output_k_cache = [] - output_v_cache = [] - # following tensors should be invariant across batches - freqs_cos = ( - self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos - ) - freqs_sin = ( - self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin - ) - - hidden_states = self.embedding_scale_factor * hidden_states - - for ind, decoder_layer in enumerate(self.layers): - k_caches = None - v_caches = None - if self.use_kv_cache: - offset_k = ind - offset_v = self.n_layers + offset_k - k_caches = args[offset_k] - v_caches = args[offset_v] - - hidden_states, k, v = decoder_layer( - hidden_states, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - atten_mask=atten_mask, - k_caches=k_caches, - v_caches=v_caches, - ) - output_k_cache.extend(k) - output_v_cache.extend(v) - - hidden_states = self.norm(hidden_states) - logits = self.output(hidden_states) - - if self.logits_scaling: - logits = logits / self.logits_scaling - - if self.output_cache: - return logits, output_k_cache, output_v_cache - return logits - - def get_example_inputs(self): - hidden_states = torch.randn( - (self.max_batch_size, self.ar_len, self.dim), dtype=torch.float32 - ) - atten_mask = AttentionMask( - CausalAttentionMask(self.max_batch_size, self.ar_len, self.max_seq_len) - ) - if self.use_kv_cache: - pos_ids = torch.zeros((self.max_batch_size, self.ar_len), dtype=torch.int32) - k_cache, v_cache = [], [] - - for _ in range(self.n_layers): - # transpose first to decrease the runtime efforts - k_cache.append( - torch.zeros( - self.max_batch_size, - self.n_kv_heads, - self.head_dim, - self.max_seq_len - self.ar_len, - ) - ) - v_cache.append( - torch.zeros( - self.max_batch_size, - self.n_kv_heads, - self.max_seq_len - self.ar_len, - self.head_dim, - ) - ) - return ( - hidden_states, - atten_mask, - pos_ids, - k_cache, - v_cache, - ) - - return ( - hidden_states, - atten_mask, - ) - - def get_metadata(self): - meta_data = super().get_metadata() - if self.audio_token_id: - meta_data["audio_token_id"] = self.audio_token_id - if self.image_token_id: - meta_data["image_token_id"] = self.image_token_id - return meta_data - - -class MultiScopeAwareLlamaModel(LlamaModel): - def __init__( - self, - config: ModelArgs, - ar_len=1, - output_new_cache_only=True, - output_cache=True, - use_i64_token=False, - **kwargs, - ): - super().__init__( - config=config, - ar_len=ar_len, - output_new_cache_only=output_new_cache_only, - output_cache=output_cache, - use_i64_token=use_i64_token, - **kwargs, - ) - # Parameter final_logit_softcapping is not necessary for all - self.final_logit_softcapping = config.final_logit_softcapping - - # Gemma2/Gemma3 requires additional configuration parameters: - # - layer_types: Specifies the type of each layer (e.g., full vs. sliding attention) - # - local_rope_theta: Base frequency for local RoPE - # - sliding_window: Size of the sliding window for local attention - self.layer_types = config.layer_types - if self.layer_types is not None: - assert len(self.layer_types) == self.n_layers, ( - f"Length of layer_types ({len(self.layer_types)}) must match " - f"n_layers ({self.n_layers})" - ) - assert ( - config.local_rope_theta is not None - ), "local_rope_theta should not be None, please set it explicitly in config." - - self.sliding_window = config.sliding_window - local_freqs_cos, local_freqs_sin = hf_precompute_freqs_cis( - config.head_dim, - config.max_context_len, - config.local_rope_theta, - config.partial_rotary_factor, - ) - local_freqs_cos = local_freqs_cos[:, : local_freqs_cos.shape[-1] // 2] - local_freqs_sin = local_freqs_sin[:, : local_freqs_sin.shape[-1] // 2] - self.register_buffer("local_freqs_cos", local_freqs_cos, persistent=False) - self.register_buffer("local_freqs_sin", local_freqs_sin, persistent=False) - - def forward( - self, - tokens: torch.Tensor, - atten_mask: torch.Tensor, - window_atten_mask: torch.Tensor, - input_pos: Optional[torch.Tensor] = None, - *args, - ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: - - output_k_cache = [] - output_v_cache = [] - # following tensors should be invariant across batches - freqs_cos = ( - self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos - ) - freqs_sin = ( - self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin - ) - local_freqs_cos = ( - self.local_freqs_cos[input_pos][0] - if self.use_kv_cache - else self.local_freqs_cos - ) - local_freqs_sin = ( - self.local_freqs_sin[input_pos][0] - if self.use_kv_cache - else self.local_freqs_sin - ) - - hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) - for ind, decoder_layer in enumerate(self.layers): - k_caches = None - v_caches = None - if self.use_kv_cache: - offset_k = ind - offset_v = self.n_layers + offset_k - k_caches = args[offset_k] - v_caches = args[offset_v] - - if self.layer_types[ind] == "sliding_attention": - hidden_states, k, v = decoder_layer( - hidden_states, - freqs_cos=local_freqs_cos, - freqs_sin=local_freqs_sin, - atten_mask=window_atten_mask, - k_caches=k_caches, - v_caches=v_caches, - ) - else: - hidden_states, k, v = decoder_layer( - hidden_states, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - atten_mask=atten_mask, - k_caches=k_caches, - v_caches=v_caches, - ) - - output_k_cache.extend(k) - output_v_cache.extend(v) - - hidden_states = self.norm(hidden_states) - logits = self.output(hidden_states) - if self.final_logit_softcapping: - logits = logits / self.final_logit_softcapping - logits = torch.tanh(logits) - logits = logits * self.final_logit_softcapping - - if self.output_cache: - return logits, output_k_cache, output_v_cache - return logits - - def get_example_inputs(self): - inputs = list(super().get_example_inputs()) - causal_mask = CausalAttentionMask( - self.max_batch_size, self.ar_len, self.max_context_len - ) - sliding_window_mask = SlidingWindowAttentionMask( - self.max_batch_size, - self.ar_len, - self.max_context_len, - sliding_window=self.sliding_window, - ) - # Don't reverse the order of attention mask - inputs[1] = AttentionMask([causal_mask, sliding_window_mask]) - return tuple(inputs) - - def get_metadata(self): - meta_data = super().get_metadata() - meta_data["get_sliding_window"] = self.sliding_window - return meta_data diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/__init__.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/__init__.py new file mode 100644 index 00000000000..e8db87faf39 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/__init__.py @@ -0,0 +1,72 @@ +# 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. + +from .attention_sink import AttentionSinkRope +from .blocks import ( + ATTENTION_REGISTRY, + DECODER_LAYER_REGISTRY, + FEED_FORWARD_REGISTRY, + FeedForward, + FeedForwardBase, + Gemma4Attention, + LayerNorm, + LlamaAttention, + LlamaDecoderLayer, + Norm, + NORM_REGISTRY, + register_attention, + register_decoder_layer, + register_feed_forward, + register_norm, + repeat_kv, + RMSNorm, + YOCOCrossDecoderLayer, +) +from .embedding import TokenEmbedding +from .rope import ( + register_rope, + register_rope_precompute, + ROPE_PRECOMPUTE_REGISTRY, + ROPE_REGISTRY, + RopeFreqs, +) +from .static_llama import ( + LlamaModel, + LlamaModelWithoutEmbedding, + MultiScopeAwareLlamaModel, +) + + +__all__ = [ + "AttentionSinkRope", + "FEED_FORWARD_REGISTRY", + "FeedForward", + "FeedForwardBase", + "register_feed_forward", + "NORM_REGISTRY", + "Norm", + "LayerNorm", + "RMSNorm", + "register_norm", + "ROPE_REGISTRY", + "ROPE_PRECOMPUTE_REGISTRY", + "RopeFreqs", + "register_rope", + "register_rope_precompute", + "repeat_kv", + "LlamaAttention", + "LlamaDecoderLayer", + "TokenEmbedding", + "LlamaModel", + "LlamaModelWithoutEmbedding", + "MultiScopeAwareLlamaModel", + "Gemma4Attention", + "ATTENTION_REGISTRY", + "register_attention", + "YOCOCrossDecoderLayer", + "DECODER_LAYER_REGISTRY", + "register_decoder_layer", +] diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/attention_sink.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/attention_sink.py new file mode 100644 index 00000000000..1c569afb548 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/attention_sink.py @@ -0,0 +1,246 @@ +# 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. + +# TODO: reenable pyre after fixing the issues +# pyre-ignore-all-errors + +import math +from typing import List + +import scipy +import torch +import torch.nn as nn + +from executorch.examples.models.llama.model_args import ModelArgs + +from executorch.examples.models.llama.rope import ( + hf_precompute_freqs_cis, + precompute_freqs_cis, +) + +from .rope import ROPE_REGISTRY + + +class AttentionSinkRope(nn.Module): + def __init__( + self, + config: ModelArgs, + sink_size: int, + eviction_batch_size: int, + ar_len: int, + **kwargs, + ): + super().__init__() + self.config = config + self.sink_size = sink_size + self.eviction_batch_size = eviction_batch_size + self.n_layers = config.n_layers + self.original_position = eviction_batch_size + sink_size + self.new_position = sink_size + self.num_to_keep = ( + config.max_context_len - sink_size - eviction_batch_size - ar_len + ) + self.evict_k_cache_shape = [ + config.max_batch_size, + config.n_kv_heads, + config.head_dim, + eviction_batch_size, + ] + self.evict_v_cache_shape = [ + config.max_batch_size, + config.n_kv_heads, + eviction_batch_size, + config.head_dim, + ] + self.kv_cache_shape = { + # single head, k input + "k": ( + config.max_batch_size, + config.n_kv_heads, + config.head_dim, + config.max_context_len - ar_len, + ), + # single head, v input + "v": ( + config.max_batch_size, + config.n_kv_heads, + config.max_context_len - ar_len, + config.head_dim, + ), + } + + if getattr(config, "enable_r3", False): + self.register_buffer( + "r3_weight", + torch.tensor( + scipy.linalg.hadamard(config.head_dim, dtype=float) + / math.sqrt(config.head_dim), + dtype=torch.float32, + device="cpu", + ), + persistent=False, + ) + + if config.partial_rotary_factor < 1: + self.apply_rope_emb = ROPE_REGISTRY["partial"] + else: + self.apply_rope_emb = ROPE_REGISTRY["default"] + + if config.use_hf_rope: + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.rope_freq_base, + config.partial_rotary_factor, + ) + freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] + freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] + else: + freqs_cos, freqs_sin = precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.rope_freq_base, + config.use_scaled_rope, + config.rope_scale_factor, + ) + original_freqs_cos = freqs_cos.narrow( + 0, self.original_position, self.num_to_keep + ) + original_freqs_sin = freqs_sin.narrow( + 0, self.original_position, self.num_to_keep + ) + new_freqs_cos = freqs_cos.narrow(0, self.new_position, self.num_to_keep) + new_freqs_sin = freqs_sin.narrow(0, self.new_position, self.num_to_keep) + rerotation_cos = ( + new_freqs_cos * original_freqs_cos + new_freqs_sin * original_freqs_sin + ) + rerotation_sin = ( + new_freqs_sin * original_freqs_cos - new_freqs_cos * original_freqs_sin + ) + self.register_buffer("rerotation_cos", rerotation_cos, persistent=False) + self.register_buffer("rerotation_sin", rerotation_sin, persistent=False) + + self.sliding_window = kwargs.get("sliding_window", False) + if self.sliding_window: + # Get attention type for each layer + self.layer_types = kwargs["layer_types"] + # Get local freq base for sliding attention + rope_freq_base = kwargs["rope_local_base_freq"] + local_freqs_cos, local_freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_context_len, + rope_freq_base, + config.partial_rotary_factor, + ) + local_freqs_cos = local_freqs_cos[:, : local_freqs_cos.shape[-1] // 2] + local_freqs_sin = local_freqs_sin[:, : local_freqs_sin.shape[-1] // 2] + local_original_freqs_cos = local_freqs_cos.narrow( + 0, self.original_position, self.num_to_keep + ) + local_original_freqs_sin = local_freqs_sin.narrow( + 0, self.original_position, self.num_to_keep + ) + local_new_freqs_cos = local_freqs_cos.narrow( + 0, self.new_position, self.num_to_keep + ) + local_new_freqs_sin = local_freqs_sin.narrow( + 0, self.new_position, self.num_to_keep + ) + local_rerotation_cos = ( + local_new_freqs_cos * local_original_freqs_cos + + local_new_freqs_sin * local_original_freqs_sin + ) + local_rerotation_sin = ( + local_new_freqs_sin * local_original_freqs_cos + - local_new_freqs_cos * local_original_freqs_sin + ) + self.register_buffer( + "local_rerotation_cos", local_rerotation_cos, persistent=False + ) + self.register_buffer( + "local_rerotation_sin", local_rerotation_sin, persistent=False + ) + + def forward(self, k_caches: List[torch.Tensor], v_caches: List[torch.Tensor]): + """ + Rerotate k_cache from original_position to new_position, and return the kv cache after eviction. This is done by rerotating + k_cache with (new_position * theta - original_position * theta) with the following matrix: + (cos(delta), -sin(delta) + sin(delta), cos(delta)) + where delta = new_position * theta - original_position * theta + + Based on https://github.com/huggingface/transformers/pull/26681 + """ + + output_k_caches, output_v_caches = [], [] + for ind, (k_cache, v_cache) in enumerate(zip(k_caches, v_caches)): + # k_cache: (batch_size, n_kv_heads, head_dim, seq_len) + # v_cache: (batch_size, n_kv_heads, seq_len, head_dim) + k_dim_to_slice = 3 + v_dim_to_slice = 2 + + k_to_keep = k_cache.narrow( + k_dim_to_slice, + self.original_position, + self.num_to_keep, + ) + k_to_keep = k_to_keep.transpose(2, 3) + if getattr(self.config, "enable_r3", False): + # We need to revert the key from spin quant before applying RoPE + k_to_keep = torch.matmul(k_to_keep, self.r3_weight.T) + + if self.sliding_window and self.layer_types[ind] == "sliding_attention": + k_to_keep = self.apply_rope_emb( + k_to_keep, self.local_rerotation_cos, self.local_rerotation_sin + ) + else: + k_to_keep = self.apply_rope_emb( + k_to_keep, self.rerotation_cos, self.rerotation_sin + ) + if getattr(self.config, "enable_r3", False): + k_to_keep = torch.matmul(k_to_keep, self.r3_weight) + k_to_keep = k_to_keep.transpose(2, 3) + new_k_cache = torch.cat( + [ + k_cache.narrow(k_dim_to_slice, 0, self.sink_size), + k_to_keep, + torch.zeros(self.evict_k_cache_shape), + ], + dim=k_dim_to_slice, + ) + + new_v_cache = torch.cat( + [ + v_cache.narrow(v_dim_to_slice, 0, self.sink_size), + v_cache.narrow( + v_dim_to_slice, + self.original_position, + self.num_to_keep, + ), + torch.zeros(self.evict_v_cache_shape), + ], + dim=v_dim_to_slice, + ) + + output_k_caches.append(new_k_cache) + output_v_caches.append(new_v_cache) + + return output_k_caches, output_v_caches + + def get_example_inputs(self): + k_cache, v_cache = [], [] + + for _ in range(self.n_layers): + k_cache.append(torch.zeros(self.kv_cache_shape["k"])) + v_cache.append(torch.zeros(self.kv_cache_shape["v"])) + return k_cache, v_cache + + def get_metadata(self): + return { + "get_eviction_batch_size": self.eviction_batch_size, + "get_max_context_len": self.config.max_context_len, + "get_sink_size": self.sink_size, + } diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/__init__.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/__init__.py new file mode 100644 index 00000000000..a21eda9f889 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/__init__.py @@ -0,0 +1,48 @@ +# 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. + +from .attention import ( + ATTENTION_REGISTRY, + Gemma4Attention, + LlamaAttention, + register_attention, + repeat_kv, +) +from .decoder_layer import ( + DECODER_LAYER_REGISTRY, + LlamaDecoderLayer, + register_decoder_layer, + YOCOCrossDecoderLayer, +) +from .feed_forward import ( + FEED_FORWARD_REGISTRY, + FeedForward, + FeedForwardBase, + register_feed_forward, +) +from .norm import LayerNorm, Norm, NORM_REGISTRY, register_norm, RMSNorm + + +__all__ = [ + "LlamaAttention", + "Gemma4Attention", + "ATTENTION_REGISTRY", + "register_attention", + "repeat_kv", + "LlamaDecoderLayer", + "YOCOCrossDecoderLayer", + "DECODER_LAYER_REGISTRY", + "register_decoder_layer", + "FEED_FORWARD_REGISTRY", + "FeedForward", + "FeedForwardBase", + "register_feed_forward", + "NORM_REGISTRY", + "Norm", + "LayerNorm", + "RMSNorm", + "register_norm", +] diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/attention.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/attention.py new file mode 100644 index 00000000000..4dcff1c7c87 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/attention.py @@ -0,0 +1,549 @@ +# 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. + +# TODO: reenable pyre after fixing the issues +# pyre-ignore-all-errors + +import math +from typing import Dict, List, Optional, Tuple, Type + +import scipy +import torch +import torch.nn as nn + +from executorch.examples.models.llama.model_args import ModelArgs + +from ..rope import ROPE_REGISTRY + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +ATTENTION_REGISTRY: Dict[str, Type[nn.Module]] = {} + + +def register_attention(model_architecture: str): + """Register an attention class by model_architecture, mirroring + FEED_FORWARD_REGISTRY: LlamaAttention is the default fallback (not + registered), while gemma4 registers its own class, which internally derives + self vs cross (KV-shared) behaviour from config.is_kv_shared_layer(layer_idx).""" + + def decorator(cls: Type[nn.Module]): + ATTENTION_REGISTRY[model_architecture] = cls + return cls + + return decorator + + +class LlamaAttention(nn.Module): + def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): + super().__init__() + self.config = config + self.dim = config.dim + self.n_heads = config.n_heads + self.head_dim = config.head_dim + self.n_kv_heads = config.n_kv_heads + self.num_key_value_groups = config.n_heads // self.n_kv_heads + self.use_kv_cache = config.use_kv_cache + self.output_new_cache_only = output_new_cache_only + self.enable_masked_softmax = getattr(config, "enable_masked_softmax", False) + self.use_qk_norm = config.use_qk_norm + self.qk_norm_before_rope = config.qk_norm_before_rope + # If None, assume each layer uses rope + self.use_rope = ( + config.no_rope_layer_interval is None + or (layer_idx + 1) % config.no_rope_layer_interval + ) + + if self.use_qk_norm: + q_norm_dim = self.head_dim + k_norm_dim = self.head_dim + self.q_norm_fn = torch.nn.RMSNorm(q_norm_dim, eps=config.norm_eps) + self.k_norm_fn = torch.nn.RMSNorm(k_norm_dim, eps=config.norm_eps) + + if config.partial_rotary_factor < 1: + self.apply_rope_emb = ROPE_REGISTRY["partial"] + else: + self.apply_rope_emb = ROPE_REGISTRY["default"] + + self.wq = nn.Linear( + self.dim, + self.n_heads * self.head_dim, + bias=config.attention_qkv_bias, + ) + self.wk = nn.Linear( + self.dim, + self.n_kv_heads * self.head_dim, + bias=config.attention_qkv_bias, + ) + self.wv = nn.Linear( + self.dim, + self.n_kv_heads * self.head_dim, + bias=config.attention_qkv_bias, + ) + self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) + + self.attn_softmax = torch.nn.Softmax(dim=-1) + + self.scale = ( + float(self.head_dim) ** 0.5 + if config.attention_multiplier is None + else 1.0 / config.attention_multiplier + ) + + # gemma 2 uses soft-capping on attention and logits + self.attn_logit_softcapping = config.attn_logit_softcapping + + if getattr(config, "enable_r3", False): + self.register_buffer( + "r3_weight", + torch.tensor( + scipy.linalg.hadamard(self.head_dim, dtype=float) + / math.sqrt(self.head_dim), + dtype=torch.float32, + device="cpu", + ), + persistent=False, + ) + + def prepare_attention_conv(self): + self.wq_conv = nn.Conv2d( + self.dim, + self.n_heads * self.head_dim, + 1, + bias=self.config.attention_qkv_bias, + ) + self.wk_conv = nn.Conv2d( + self.dim, + self.n_kv_heads * self.head_dim, + 1, + bias=self.config.attention_qkv_bias, + ) + self.wv_conv = nn.Conv2d( + self.dim, + self.n_kv_heads * self.head_dim, + 1, + bias=self.config.attention_qkv_bias, + ) + self.wo_conv = nn.Conv2d(self.n_heads * self.head_dim, self.dim, 1, bias=False) + + self.forward_no_conv = self.forward + self.forward = self.forward_attention_conv + + self.wq_conv.weight.data.copy_(self.wq.weight[:, :, None, None]) + if self.wq_conv.bias is not None: + self.wq_conv.bias.data.copy_(self.wq.bias) + self.wk_conv.weight.data.copy_(self.wk.weight[:, :, None, None]) + if self.wk_conv.bias is not None: + self.wk_conv.bias.data.copy_(self.wk.bias) + self.wv_conv.weight.data.copy_(self.wv.weight[:, :, None, None]) + if self.wv_conv.bias is not None: + self.wv_conv.bias.data.copy_(self.wv.bias) + self.wo_conv.weight.data.copy_(self.wo.weight[:, :, None, None]) + + del self.wq + del self.wk + del self.wv + del self.wo + + def forward_attention_conv( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + k_caches: List[torch.Tensor], + v_caches: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, seq_len, _ = hidden_states.shape + hidden_states = torch.reshape( + hidden_states, (bsz, seq_len, 1, self.dim) + ).transpose(1, 3) + + q = self.wq_conv(hidden_states) + k = self.wk_conv(hidden_states) + v = self.wv_conv(hidden_states) + q = q.permute(0, 3, 1, 2).squeeze(-1) + k = k.permute(0, 3, 1, 2).squeeze(-1) + v = v.permute(0, 3, 1, 2).squeeze(-1) + q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + + if self.use_qk_norm and self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) + + if self.use_rope: + q = self.apply_rope_emb(q, freqs_cos, freqs_sin) + k = self.apply_rope_emb(k, freqs_cos, freqs_sin) + + if self.use_qk_norm and not self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) + if getattr(self.config, "enable_r3", False): + q = torch.matmul(q, self.r3_weight) + k = torch.matmul(k, self.r3_weight) + k = k.transpose(2, 3) + + kh, vh = None, None + # kv cache mode + if self.use_kv_cache: + kh = torch.cat([k_caches, k], dim=-1) + vh = torch.cat([v_caches, v], dim=2) + # batch_prefill mode + else: + kh = k + vh = v + + kh = repeat_kv(kh, self.num_key_value_groups) + vh = repeat_kv(vh, self.num_key_value_groups) + + attn = q @ kh + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2) + y = y.reshape(bsz, seq_len, 1, -1).transpose(1, 3) + y = self.wo_conv(y) + y = y.transpose(1, 3) + y = y.reshape(bsz, seq_len, -1) + + if self.output_new_cache_only: + return y, [k], [v] + + return y, [kh], [vh] + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + k_caches: List[torch.Tensor], + v_caches: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, seq_len, _ = hidden_states.shape + + q = self.wq(hidden_states) + k = self.wk(hidden_states) + v = self.wv(hidden_states) + q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + + if self.use_qk_norm and self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) + + if self.use_rope: + q = self.apply_rope_emb(q, freqs_cos, freqs_sin) + k = self.apply_rope_emb(k, freqs_cos, freqs_sin) + + if self.use_qk_norm and not self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) + if getattr(self.config, "enable_r3", False): + q = torch.matmul(q, self.r3_weight) + k = torch.matmul(k, self.r3_weight) + k = k.transpose(2, 3) + + kh, vh = None, None + # kv cache mode + if self.use_kv_cache: + kh = torch.cat([k_caches, k], dim=-1) + vh = torch.cat([v_caches, v], dim=2) + # batch_prefill mode + else: + kh = k + vh = v + + kh = repeat_kv(kh, self.num_key_value_groups) + vh = repeat_kv(vh, self.num_key_value_groups) + + attn = q @ kh + # gemma2-2b + if self.attn_logit_softcapping is not None: + attn = attn / self.attn_logit_softcapping + attn = torch.tanh(attn) + attn = attn * self.attn_logit_softcapping + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2).reshape(bsz, seq_len, -1) + y = self.wo(y) + + if self.output_new_cache_only: + return y, [k], [v] + + return y, [kh], [vh] + + +@register_attention("gemma4") +class Gemma4Attention(nn.Module): + """Gemma4 attention. Self layers compute their own K/V (and donor layers + expose them for YOCO sharing); cross layers (KV-shared) reuse a donor + layer's K/V and produce no new cache. All layers keep the K/V projections + (checkpoint parity); config.is_kv_shared_layer(layer_idx) only changes the + forward path, not the parameter set.""" + + def __init__( + self, + layer_idx: int, + config: ModelArgs, + output_new_cache_only: bool = True, + ): + super().__init__() + self.config = config + self.dim = config.dim + self.n_heads = config.n_heads + self.n_kv_heads = config.n_kv_heads + self.num_key_value_groups = config.n_heads // self.n_kv_heads + self.use_kv_cache = config.use_kv_cache + self.output_new_cache_only = output_new_cache_only + self.enable_masked_softmax = getattr(config, "enable_masked_softmax", False) + self.is_kv_shared_layer = config.is_kv_shared_layer(layer_idx) + + # Full-attention layers use global_head_dim; sliding layers use head_dim. + self.head_dim = config.get_head_dim(layer_idx) + + self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False) + self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) + self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) + self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) + # Gemma4 always applies QK-norm (before rope) and V-norm. KV-shared + # (cross) layers keep the K/V projections for checkpoint parity even + # though their forward reuses a donor's K/V instead. + self.q_norm_fn = torch.nn.RMSNorm(self.head_dim, eps=config.norm_eps) + self.k_norm_fn = torch.nn.RMSNorm(self.head_dim, eps=config.norm_eps) + self.v_norm_fn = torch.nn.RMSNorm( + self.head_dim, eps=config.norm_eps, elementwise_affine=False + ) + + # Sliding layers apply full rope; full-attention layers use partial rope. + if config.is_sliding_attention(layer_idx): + self.apply_rope_emb = ROPE_REGISTRY["default"] + else: + self.apply_rope_emb = ROPE_REGISTRY["partial"] + + self.attn_softmax = torch.nn.Softmax(dim=-1) + + # Gemma 4 uses scale=1.0; QK-norm handles normalisation. + self.scale = 1.0 + + if getattr(config, "enable_r3", False): + self.register_buffer( + "r3_weight", + torch.tensor( + scipy.linalg.hadamard(self.head_dim, dtype=float) + / math.sqrt(self.head_dim), + dtype=torch.float32, + device="cpu", + ), + persistent=False, + ) + + def prepare_attention_conv(self): + self.wq_conv = nn.Conv2d(self.dim, self.n_heads * self.head_dim, 1, bias=False) + self.wo_conv = nn.Conv2d(self.n_heads * self.head_dim, self.dim, 1, bias=False) + self.wk_conv = nn.Conv2d( + self.dim, self.n_kv_heads * self.head_dim, 1, bias=False + ) + self.wv_conv = nn.Conv2d( + self.dim, self.n_kv_heads * self.head_dim, 1, bias=False + ) + self.wq_conv.weight.data.copy_(self.wq.weight[:, :, None, None]) + self.wo_conv.weight.data.copy_(self.wo.weight[:, :, None, None]) + self.wk_conv.weight.data.copy_(self.wk.weight[:, :, None, None]) + self.wv_conv.weight.data.copy_(self.wv.weight[:, :, None, None]) + del self.wq + del self.wo + del self.wk + del self.wv + + self.forward_no_conv = self.forward + self.forward = self.forward_attention_conv + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + k_caches: Optional[torch.Tensor] = None, + v_caches: Optional[torch.Tensor] = None, + donor_k: Optional[torch.Tensor] = None, + donor_v: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + bsz, seq_len, _ = hidden_states.shape + + q = self.wq(hidden_states) + q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) + q = self.q_norm_fn(q) + q = self.apply_rope_emb(q, freqs_cos, freqs_sin) + if getattr(self.config, "enable_r3", False): + q = torch.matmul(q, self.r3_weight) + + if self.is_kv_shared_layer: + kh = repeat_kv(donor_k, self.num_key_value_groups) + vh = repeat_kv(donor_v, self.num_key_value_groups) + attn = q @ kh + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2).reshape(bsz, seq_len, -1) + return self.wo(y), None, None + + k = self.wk(hidden_states) + v = self.wv(hidden_states) + k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + k = self.k_norm_fn(k) + v = self.v_norm_fn(v) + k = self.apply_rope_emb(k, freqs_cos, freqs_sin) + if getattr(self.config, "enable_r3", False): + k = torch.matmul(k, self.r3_weight) + k = k.transpose(2, 3) + + if k_caches is not None: + kh = torch.cat([k_caches, k], dim=-1) + vh = torch.cat([v_caches, v], dim=2) + else: + kh = k + vh = v + + new_k = k if self.output_new_cache_only else kh + new_v = v if self.output_new_cache_only else vh + + kh = repeat_kv(kh, self.num_key_value_groups) + vh = repeat_kv(vh, self.num_key_value_groups) + attn = q @ kh + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2).reshape(bsz, seq_len, -1) + return self.wo(y), new_k, new_v + + def forward_attention_conv( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + k_caches: Optional[torch.Tensor] = None, + v_caches: Optional[torch.Tensor] = None, + donor_k: Optional[torch.Tensor] = None, + donor_v: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + bsz, seq_len, _ = hidden_states.shape + hidden_states = torch.reshape( + hidden_states, (bsz, seq_len, 1, self.dim) + ).transpose(1, 3) + + q = self.wq_conv(hidden_states) + q = q.permute(0, 3, 1, 2).squeeze(-1) + q = q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) + q = self.q_norm_fn(q) + q = self.apply_rope_emb(q, freqs_cos, freqs_sin) + if getattr(self.config, "enable_r3", False): + q = torch.matmul(q, self.r3_weight) + + if self.is_kv_shared_layer: + kh = repeat_kv(donor_k, self.num_key_value_groups) + vh = repeat_kv(donor_v, self.num_key_value_groups) + attn = q @ kh + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2) + y = y.reshape(bsz, seq_len, 1, -1).transpose(1, 3) + y = self.wo_conv(y) + y = y.transpose(1, 3) + y = y.reshape(bsz, seq_len, -1) + return y, None, None + + k = self.wk_conv(hidden_states) + v = self.wv_conv(hidden_states) + k = k.permute(0, 3, 1, 2).squeeze(-1) + v = v.permute(0, 3, 1, 2).squeeze(-1) + k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + k = self.k_norm_fn(k) + v = self.v_norm_fn(v) + k = self.apply_rope_emb(k, freqs_cos, freqs_sin) + if getattr(self.config, "enable_r3", False): + k = torch.matmul(k, self.r3_weight) + k = k.transpose(2, 3) + + if k_caches is not None: + kh = torch.cat([k_caches, k], dim=-1) + vh = torch.cat([v_caches, v], dim=2) + else: + kh = k + vh = v + + new_k = k if self.output_new_cache_only else kh + new_v = v if self.output_new_cache_only else vh + + kh = repeat_kv(kh, self.num_key_value_groups) + vh = repeat_kv(vh, self.num_key_value_groups) + attn = q @ kh + attn = attn / self.scale + if self.enable_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where(atten_mask == 0, attn, attn_min + minus_value) + else: + attn = attn + atten_mask + attn = self.attn_softmax(attn) + y = attn @ vh + y = y.transpose(1, 2) + y = y.reshape(bsz, seq_len, 1, -1).transpose(1, 3) + y = self.wo_conv(y) + y = y.transpose(1, 3) + y = y.reshape(bsz, seq_len, -1) + return y, new_k, new_v diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/decoder_layer.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/decoder_layer.py new file mode 100644 index 00000000000..cb5bbc0625b --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/decoder_layer.py @@ -0,0 +1,183 @@ +# 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. + +# TODO: reenable pyre after fixing the issues +# pyre-ignore-all-errors + +from typing import Dict, List, Optional, Tuple, Type + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from executorch.examples.models.llama.model_args import ModelArgs + +from .attention import ATTENTION_REGISTRY, LlamaAttention +from .feed_forward import FEED_FORWARD_REGISTRY, FeedForward +from .norm import NORM_REGISTRY + + +DECODER_LAYER_REGISTRY: Dict[bool, Type[nn.Module]] = {} + + +def register_decoder_layer(is_kv_shared: bool): + def decorator(cls: Type[nn.Module]): + DECODER_LAYER_REGISTRY[is_kv_shared] = cls + return cls + + return decorator + + +class LlamaDecoderLayer(nn.Module): + + def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): + super().__init__() + self.dim = config.dim + self.attention = ATTENTION_REGISTRY.get( + config.model_architecture, LlamaAttention + )( + layer_idx=layer_idx, + config=config, + output_new_cache_only=output_new_cache_only, + ) + + self.feed_forward = FEED_FORWARD_REGISTRY.get( + config.model_architecture, FeedForward + )(config, hidden_dim=config.get_intermediate_size(layer_idx)) + self.attention_norm = NORM_REGISTRY[config.norm_type]( + config.dim, eps=config.norm_eps + ) + self.ffn_norm = ( + NORM_REGISTRY[config.norm_type](config.dim, eps=config.norm_eps) + if config.use_ffn_norm + else None + ) + self.post_attention_norm = ( + torch.nn.RMSNorm(config.dim, eps=config.norm_eps) + if config.post_attention_norm + else None + ) + self.residual_multiplier = config.residual_multiplier + self.post_ffn_norm = ( + torch.nn.RMSNorm(config.dim, eps=config.norm_eps) + if config.post_ffn_norm + else None + ) + + # Per-Layer Embeddings (Currently only used by Gemma4). + self.use_per_layer_embedding = ( + config.vocab_size_per_layer_input and config.hidden_size_per_layer_input + ) + if self.use_per_layer_embedding: + self.per_layer_input_gate = nn.Linear( + config.dim, config.hidden_size_per_layer_input, bias=False + ) + self.per_layer_projection = nn.Linear( + config.hidden_size_per_layer_input, config.dim, bias=False + ) + self.layer_scalar = nn.Parameter(torch.ones(1)) + self.post_per_layer_input_norm = NORM_REGISTRY[config.norm_type]( + config.dim, eps=config.norm_eps + ) + + def forward( + self, + x: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + k_caches: List[torch.Tensor], + v_caches: List[torch.Tensor], + per_layer_input: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_states = self.attention_norm(x) + h, k_cache, v_cache = self.attention( + hidden_states, + freqs_cos, + freqs_sin, + atten_mask, + k_caches, + v_caches, + ) + if self.post_attention_norm: + h = self.post_attention_norm(h) + h = ( + x + h * self.residual_multiplier + if self.residual_multiplier is not None + else x + h + ) + hidden_states = hidden_states if self.ffn_norm is None else self.ffn_norm(h) + out = self.feed_forward(hidden_states) + if self.post_ffn_norm: + out = self.post_ffn_norm(out) + output = ( + h + out * self.residual_multiplier + if self.residual_multiplier is not None + else h + out + ) + + # PLE + if self.use_per_layer_embedding: + gated = F.gelu(self.per_layer_input_gate(output), approximate="tanh") + projected = self.per_layer_projection(gated * per_layer_input) + output = ( + output + self.post_per_layer_input_norm(projected) + ) * self.layer_scalar + + return output, k_cache, v_cache + + +@register_decoder_layer(is_kv_shared=True) +class YOCOCrossDecoderLayer(LlamaDecoderLayer): + """Cross-decoder layer (YOCO): attention reuses a donor layer's K/V and + produces no new cache. Shares the self layer's submodules and residual / + FFN / PLE flow, overriding only the attention call.""" + + def forward( + self, + x: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + atten_mask: torch.Tensor, + donor_k: Optional[torch.Tensor], + donor_v: Optional[torch.Tensor], + per_layer_input: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, None, None]: + hidden_states = self.attention_norm(x) + h, _, _ = self.attention( + hidden_states, + freqs_cos, + freqs_sin, + atten_mask, + donor_k=donor_k, + donor_v=donor_v, + ) + if self.post_attention_norm: + h = self.post_attention_norm(h) + h = ( + x + h * self.residual_multiplier + if self.residual_multiplier is not None + else x + h + ) + hidden_states = hidden_states if self.ffn_norm is None else self.ffn_norm(h) + out = self.feed_forward(hidden_states) + if self.post_ffn_norm: + out = self.post_ffn_norm(out) + output = ( + h + out * self.residual_multiplier + if self.residual_multiplier is not None + else h + out + ) + + # PLE + if self.use_per_layer_embedding: + gated = F.gelu(self.per_layer_input_gate(output), approximate="tanh") + projected = self.per_layer_projection(gated * per_layer_input) + output = ( + output + self.post_per_layer_input_norm(projected) + ) * self.layer_scalar + + return output, None, None diff --git a/examples/qualcomm/oss_scripts/llama/model/feed_forward.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/feed_forward.py similarity index 67% rename from examples/qualcomm/oss_scripts/llama/model/feed_forward.py rename to examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/feed_forward.py index 2f36779cc71..3f9ec61e5ae 100644 --- a/examples/qualcomm/oss_scripts/llama/model/feed_forward.py +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/feed_forward.py @@ -7,6 +7,8 @@ from typing import Dict, Type import torch +import torch.nn as nn + from executorch.examples.models.llama.model_args import ModelArgs from transformers.activations import GELUActivation @@ -25,24 +27,62 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: pass -FeedForward_REGISTRY: Dict[str, Type[FeedForwardBase]] = {} +FEED_FORWARD_REGISTRY: Dict[str, Type[FeedForwardBase]] = {} def register_feed_forward(name: str): """Decorator to register norm classes""" def decorator(cls: Type[FeedForwardBase]): - FeedForward_REGISTRY[name] = cls + FEED_FORWARD_REGISTRY[name] = cls return cls return decorator +class FeedForward(nn.Module): + def __init__(self, args: ModelArgs, dim=None, hidden_dim=None): + super().__init__() + self.dim: int = dim if dim is not None else args.dim + self.hidden_dim: int = hidden_dim if hidden_dim is not None else args.hidden_dim + self.w1 = nn.Linear(self.dim, self.hidden_dim, bias=False) + self.w2 = nn.Linear(self.hidden_dim, self.dim, bias=False) + self.w3 = nn.Linear(self.dim, self.hidden_dim, bias=False) + self.act_fn = args.act_fn.get_function() + + def prepare_feedforward_conv(self): + self.w1_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False) + self.w2_conv = nn.Conv2d(self.hidden_dim, self.dim, 1, bias=False) + self.w3_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False) + + self.forward_no_conv = self.forward + self.forward = self.forward_feedforward_conv + self.w1_conv.weight.data.copy_(self.w1.weight[:, :, None, None]) + self.w2_conv.weight.data.copy_(self.w2.weight[:, :, None, None]) + self.w3_conv.weight.data.copy_(self.w3.weight[:, :, None, None]) + + del self.w1 + del self.w2 + del self.w3 + + def forward_feedforward_conv(self, x): + bsz, _, _ = x.size() + x = torch.reshape(x, (bsz, -1, 1, self.dim)) + x = x.transpose(1, 3) # Transpose right before and after Conv + x = self.w2_conv(self.act_fn(self.w1_conv(x)) * self.w3_conv(x)) + x = x.transpose(1, 3) + x = torch.reshape(x, (bsz, -1, self.dim)) + return x + + def forward(self, x): + return self.w2(self.act_fn(self.w1(x)) * self.w3(x)) + + @register_feed_forward("CodeGenModel") class CodegenFeedForward(FeedForwardBase): """FeedForward with fc_in and fc_out""" - def __init__(self, args: ModelArgs): # in MLP: intermediate_size= 4 * embed_dim + def __init__(self, args: ModelArgs): super().__init__() assert args.hidden_dim is not None @@ -54,13 +94,13 @@ def __init__(self, args: ModelArgs): # in MLP: intermediate_size= 4 * embed_dim # HF uses NewGelu, however, Gelu is a fused op in QNN and can run faster self.act = GELUActivation(use_gelu_python=False) - def prepare_feedfoward_conv(self): + def prepare_feedforward_conv(self): intermediate_size = 4 * self.dim self.fc_in_conv = torch.nn.Conv2d(self.dim, intermediate_size, 1, bias=True) self.fc_out_conv = torch.nn.Conv2d(self.hidden_dim, self.dim, 1, bias=True) self.forward_no_conv = self.forward - self.forward = self.forward_feedfoward_conv + self.forward = self.forward_feedforward_conv self.fc_in_conv.weight.data.copy_(self.fc_in.weight[:, :, None, None]) self.fc_out_conv.weight.data.copy_(self.fc_out.weight[:, :, None, None]) @@ -71,7 +111,7 @@ def prepare_feedfoward_conv(self): del self.fc_in del self.fc_out - def forward_feedfoward_conv(self, x): + def forward_feedforward_conv(self, x): bsz, _, _ = x.size() x = torch.reshape(x, (bsz, -1, 1, self.dim)) @@ -94,7 +134,7 @@ def forward(self, x): class GLMFeedForward(FeedForwardBase): """FeedForward with gate_up_proj and down_proj""" - def __init__(self, args: ModelArgs): # in MLP: intermediate_size= 4 * embed_dim + def __init__(self, args: ModelArgs): super().__init__() assert args.hidden_dim is not None @@ -105,14 +145,14 @@ def __init__(self, args: ModelArgs): # in MLP: intermediate_size= 4 * embed_dim self.down_proj = torch.nn.Linear(args.hidden_dim, args.dim, bias=False) self.activation_fn = args.act_fn.get_function() - def prepare_feedfoward_conv(self): + def prepare_feedforward_conv(self): self.gate_up_proj_conv = torch.nn.Conv2d( self.dim, 2 * self.hidden_dim, 1, bias=False ) self.down_proj_conv = torch.nn.Conv2d(self.hidden_dim, self.dim, 1, bias=False) self.forward_no_conv = self.forward - self.forward = self.forward_feedfoward_conv + self.forward = self.forward_feedforward_conv self.gate_up_proj_conv.weight.data.copy_( self.gate_up_proj.weight[:, :, None, None] @@ -122,7 +162,7 @@ def prepare_feedfoward_conv(self): del self.gate_up_proj del self.down_proj - def forward_feedfoward_conv(self, x): + def forward_feedforward_conv(self, x): bsz, _, _ = x.size() x = torch.reshape(x, (bsz, -1, 1, self.dim)) x = x.transpose(1, 3) # Transpose right before and after Conv diff --git a/examples/qualcomm/oss_scripts/llama/model/layernorm.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/norm.py similarity index 100% rename from examples/qualcomm/oss_scripts/llama/model/layernorm.py rename to examples/qualcomm/oss_scripts/llama/model/text_decoder/blocks/norm.py diff --git a/examples/qualcomm/oss_scripts/llama/model/embedding.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/embedding.py similarity index 100% rename from examples/qualcomm/oss_scripts/llama/model/embedding.py rename to examples/qualcomm/oss_scripts/llama/model/text_decoder/embedding.py diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/rope.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/rope.py new file mode 100644 index 00000000000..3becc35dcb9 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/rope.py @@ -0,0 +1,186 @@ +# 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. +from dataclasses import dataclass, fields +from typing import Callable, Dict, Iterator, Optional, Tuple + +import torch + +from executorch.examples.models.llama.model_args import ModelArgs +from executorch.examples.models.llama.rope import ( + hf_precompute_freqs_cis, + precompute_freqs_cis, +) + + +ROPE_REGISTRY: Dict[str, Callable] = {} + + +def register_rope(name: str): + """Register a rotary embedding function.""" + + def decorator(fn: Callable): + ROPE_REGISTRY[name] = fn + return fn + + return decorator + + +ROPE_PRECOMPUTE_REGISTRY: Dict[str, Callable] = {} + + +def register_rope_precompute(name: str): + """Register a RopeFreqs precompute method under a precompute-type name. + + The decorated method takes only self and fills the cos/sin fields from + self.config, so callers just do RopeFreqs(config).compute().""" + + def decorator(method: Callable) -> Callable: + ROPE_PRECOMPUTE_REGISTRY[name] = method + return method + + return decorator + + +def hf_zero_pad_precompute_freqs_cis( + head_dim, max_context_len, theta, partial_rotary_factor +) -> Tuple[torch.Tensor, torch.Tensor]: + """HF-style RoPE with zero-padded non-rotated dims, so ROPE_REGISTRY["partial"] + handles the half-and-half split -- used by gemma4 on both its partial-rotation + global layers and its full-rotation local layers.""" + rotary_dim = int(head_dim * partial_rotary_factor) + half_rotary_dim = rotary_dim // 2 + inv_freq_rotary = 1.0 / ( + theta ** (torch.arange(0, rotary_dim, 2).float() / head_dim) + ) + num_no_rope = head_dim // 2 - half_rotary_dim + inv_freq = ( + torch.cat([inv_freq_rotary, torch.zeros(num_no_rope)]) + if num_no_rope > 0 + else inv_freq_rotary + ) + positions = torch.arange(max_context_len, dtype=torch.float32) + freqs = torch.outer(positions, inv_freq) + return torch.cos(freqs), torch.sin(freqs) + + +@dataclass +class RopeFreqs: + """Precomputed rope cos/sin tables plus the config that produced them. + + compute() selects the precompute method for this model and fills the cos/sin + fields. named_buffers() then yields the (field_name, tensor) pairs the model + registers directly -- the tensor field names ARE the buffer names. + Sliding-window models fill the local_* fields; others leave them None. + """ + + config: ModelArgs + freqs_cos: Optional[torch.Tensor] = None + freqs_sin: Optional[torch.Tensor] = None + local_freqs_cos: Optional[torch.Tensor] = None + local_freqs_sin: Optional[torch.Tensor] = None + + def compute(self) -> "RopeFreqs": + # gemma4 (global_head_dim set) needs zero-padded partial-rope tables so + # the buffers stay compatible with the partial apply kernel even on + # full-rotation local layers -- a case the partial_rotary_factor flag + # alone can no longer distinguish. + if self.config.global_head_dim: + precompute_type = "hf_zero_pad" + else: + precompute_type = "hf" if self.config.use_hf_rope else "scaled" + ROPE_PRECOMPUTE_REGISTRY[precompute_type](self) + return self + + def named_buffers(self) -> Iterator[Tuple[str, torch.Tensor]]: + for f in fields(self): + if f.name == "config": + continue + tensor = getattr(self, f.name) + if tensor is not None: + yield f.name, tensor + + @register_rope_precompute("scaled") + def _precompute_freqs_cis(self) -> None: + config = self.config + self.freqs_cos, self.freqs_sin = precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.rope_freq_base, + config.use_scaled_rope, + config.rope_scale_factor, + ) + + @register_rope_precompute("hf") + def _hf_precompute_freqs_cis(self) -> None: + config = self.config + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.rope_freq_base, + config.partial_rotary_factor, + ) + self.freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] + self.freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] + + if config.local_rope_theta is not None: + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.local_rope_theta, + config.partial_rotary_factor, + ) + self.local_freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] + self.local_freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] + + @register_rope_precompute("hf_zero_pad") + def _hf_zero_pad_precompute_freqs_cis(self) -> None: + """gemma4: global layers use global_head_dim + partial rope; local + (sliding) layers use head_dim + full rope, both zero-padded to match + the partial kernel.""" + config = self.config + self.freqs_cos, self.freqs_sin = hf_zero_pad_precompute_freqs_cis( + config.global_head_dim, + config.max_context_len, + config.rope_theta, + config.partial_rotary_factor, + ) + self.local_freqs_cos, self.local_freqs_sin = hf_zero_pad_precompute_freqs_cis( + config.head_dim, + config.max_context_len, + config.local_rope_theta, + 1.0, + ) + + +@register_rope("partial") +def apply_partial_rotary_emb_single(x, freqs_cos, freqs_sin): + if x.dim() == 4: + freqs_cos = freqs_cos[None, None, :, :] + freqs_sin = freqs_sin[None, None, :, :] + rotary_dim = freqs_cos.shape[-1] * 2 + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x_r, x_i = x_rot[..., : x_rot.shape[-1] // 2], x_rot[..., x_rot.shape[-1] // 2 :] + x_out_r = x_r * freqs_cos - x_i * freqs_sin + x_out_i = x_r * freqs_sin + x_i * freqs_cos + x_rotated = torch.cat([x_out_r, x_out_i], dim=-1) + return torch.cat([x_rotated, x_pass], dim=-1) + + +@register_rope("default") +def apply_rotary_emb_single(x, freqs_cos, freqs_sin): + # The implementation of RoPE in HuggingFace processes query and key with two half instead of interleaved way. + # The main difference is stride in StrideSlice op. For interleaved way, stride is two which is not friendly for HTP backend. + # Ref: https://github.com/huggingface/transformers/issues/25199 + x_r, x_i = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + # broadcast for batch_prefill mode input x + if x.dim() == 4: + freqs_cos = freqs_cos[None, None, :, :] + freqs_sin = freqs_sin[None, None, :, :] + x_out_r = x_r * freqs_cos - x_i * freqs_sin + x_out_i = x_r * freqs_sin + x_i * freqs_cos + + x_out = torch.cat([x_out_r, x_out_i], dim=-1) + return x_out diff --git a/examples/qualcomm/oss_scripts/llama/model/text_decoder/static_llama.py b/examples/qualcomm/oss_scripts/llama/model/text_decoder/static_llama.py new file mode 100644 index 00000000000..bf7cc270c8c --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/model/text_decoder/static_llama.py @@ -0,0 +1,576 @@ +# 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. + +# TODO: reenable pyre after fixing the issues +# pyre-ignore-all-errors + +import math +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn + +from executorch.examples.models.llama.model_args import ModelArgs + +from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import ( + AttentionMask, + CausalAttentionMask, + SlidingWindowAttentionMask, +) + +from .blocks.decoder_layer import DECODER_LAYER_REGISTRY, LlamaDecoderLayer +from .blocks.norm import NORM_REGISTRY + +from .rope import RopeFreqs + + +class LlamaModel(nn.Module): + def __init__( + self, + config: ModelArgs, + ar_len=1, + output_new_cache_only=True, + output_cache=True, + use_i64_token=False, + **kwargs, + ): + super().__init__() + self.dim = config.dim + self.head_dim = config.head_dim + self.max_batch_size = config.max_batch_size + self.max_seq_len = config.max_seq_len + self.max_context_len = config.max_context_len + self.n_heads = config.n_heads + self.n_kv_heads = config.n_kv_heads + self.n_layers = config.n_layers + self.n_self_layers = config.n_layers - config.num_kv_shared_layers # for YOCO + self.vocab_size = config.vocab_size + self.rope_freq_base = config.rope_freq_base + self.use_kv_cache = config.use_kv_cache + self.embedding_scale_factor = config.embedding_scale_factor + self.ar_len = ar_len + self.output_new_cache_only = output_new_cache_only + self.use_i64_token = use_i64_token + self.output_cache = output_cache + self.kv_io_bit_width = config.kv_io_bit_width + self.logits_scaling = config.logits_scaling + self.config = config + + self.layers = nn.ModuleList( + [ + DECODER_LAYER_REGISTRY.get( + config.is_kv_shared_layer(layer_idx), LlamaDecoderLayer + )( + config=config, + layer_idx=layer_idx, + output_new_cache_only=self.output_new_cache_only, + ) + for layer_idx in range(config.n_layers) + ] + ) + self.norm = NORM_REGISTRY[config.norm_type](config.dim, eps=config.norm_eps) + self.output = nn.Linear(config.dim, config.vocab_size, bias=config.output_bias) + + self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim) + for buffer_name, buffer in RopeFreqs(config).compute().named_buffers(): + self.register_buffer(buffer_name, buffer, persistent=False) + + def prepare_output_conv(self): + def forward_output_conv(x): + bsz, _, _ = x.size() + x = torch.reshape(x, (bsz, -1, 1, self.dim)) + x = x.transpose(1, 3) # Transpose right before and after Conv + x = self.output_conv(x) + x = x.transpose(1, 3) + x = torch.reshape(x, (bsz, -1, self.vocab_size)) + return x + + self.output_conv = nn.Conv2d(self.dim, self.vocab_size, 1, bias=False) + self.output_conv.weight.data.copy_(self.output.weight[:, :, None, None]) + + del self.output + self.output = forward_output_conv + + def forward( + self, + tokens: torch.Tensor, + atten_mask: torch.Tensor, + input_pos: Optional[torch.Tensor] = None, + *args, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + output_k_cache = [] + output_v_cache = [] + # following tensors should be invariant across batches + freqs_cos = ( + self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos + ) + freqs_sin = ( + self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin + ) + + hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) + + for ind, decoder_layer in enumerate(self.layers): + k_caches = None + v_caches = None + if self.use_kv_cache: + offset_k = ind + offset_v = self.n_layers + offset_k + k_caches = args[offset_k] + v_caches = args[offset_v] + + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + atten_mask=atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + output_k_cache.extend(k) + output_v_cache.extend(v) + + hidden_states = self.norm(hidden_states) + logits = self.output(hidden_states) + + if self.logits_scaling: + logits = logits / self.logits_scaling + + if self.output_cache: + return logits, output_k_cache, output_v_cache + return logits + + def get_example_inputs(self): + dtype = torch.int64 if self.use_i64_token else torch.int32 + tokens = torch.randint( + self.vocab_size, (self.max_batch_size, self.ar_len), dtype=dtype + ) + atten_mask = AttentionMask( + CausalAttentionMask(self.max_batch_size, self.ar_len, self.max_context_len) + ) + if self.use_kv_cache: + pos_ids = torch.zeros((self.max_batch_size, self.ar_len), dtype=torch.int32) + k_cache, v_cache = [], [] + for ind in range(self.n_self_layers): + head_dim = self.config.get_head_dim(ind) + # transpose first to decrease the runtime efforts + k_cache.append( + torch.zeros( + self.max_batch_size, + self.n_kv_heads, + head_dim, + self.max_context_len - self.ar_len, + ) + ) + v_cache.append( + torch.zeros( + self.max_batch_size, + self.n_kv_heads, + self.max_context_len - self.ar_len, + head_dim, + ) + ) + return ( + tokens, + atten_mask, + pos_ids, + k_cache, + v_cache, + ) + + return ( + tokens, + atten_mask, + ) + + def get_metadata(self): + return { + "get_ar_len": self.ar_len, + "get_bos_id": 1, + "get_eos_id": 2, + "get_dim": self.dim, + "get_head_dim": self.head_dim, + "get_max_batch_size": self.max_batch_size, + "get_max_seq_len": self.max_seq_len, + "get_max_context_len": self.max_context_len, + "get_n_bos": 1, + "get_n_eos": 1, + "get_n_kv_heads": self.n_kv_heads, + "get_n_layers": self.n_layers, + "get_n_self_layers": self.n_self_layers, + "get_vocab_size": self.vocab_size, + "get_use_kv_cache": self.use_kv_cache, + "get_kv_io_bit_width": self.kv_io_bit_width, + } + + def get_kv_head_dim(self): + """Head dims that appear in K/V caches -- one for most models.""" + return {self.head_dim} + + def get_kv_cache_shapes(self): + return set().union( + *( + { + # single head, kv input + (head_dim, self.max_context_len), + (self.max_context_len, head_dim), + # single head, kv output + (head_dim, self.ar_len), + (self.ar_len, head_dim), + } + for head_dim in self.get_kv_head_dim() + ) + ) + + def get_freq_shapes(self): + """The rope cos/sin (ar_len, head_dim // 2) shapes this model produces.""" + return {(self.ar_len, head_dim // 2) for head_dim in self.get_kv_head_dim()} + + +class LlamaModelWithoutEmbedding(LlamaModel): + def __init__( + self, + config: ModelArgs, + ar_len=1, + output_new_cache_only=True, + output_cache=True, + use_i64_token=False, + **kwargs, + ): + super().__init__( + config=config, + ar_len=ar_len, + output_new_cache_only=output_new_cache_only, + output_cache=output_cache, + use_i64_token=use_i64_token, + **kwargs, + ) + + # Set the audio/image token ID from keyword arguments. It defaults to None if not provided. + # If an ID is provided, it will be stored in the model's metadata. + self.audio_token_id = kwargs.get("audio_token_id", None) + self.image_token_id = kwargs.get("image_token_id", None) + + def forward( + self, + hidden_states: torch.Tensor, + atten_mask: torch.Tensor, + input_pos: Optional[torch.Tensor] = None, + *args, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + output_k_cache = [] + output_v_cache = [] + # following tensors should be invariant across batches + freqs_cos = ( + self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos + ) + freqs_sin = ( + self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin + ) + + hidden_states = self.embedding_scale_factor * hidden_states + + for ind, decoder_layer in enumerate(self.layers): + k_caches = None + v_caches = None + if self.use_kv_cache: + offset_k = ind + offset_v = self.n_layers + offset_k + k_caches = args[offset_k] + v_caches = args[offset_v] + + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + atten_mask=atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + output_k_cache.extend(k) + output_v_cache.extend(v) + + hidden_states = self.norm(hidden_states) + logits = self.output(hidden_states) + + if self.logits_scaling: + logits = logits / self.logits_scaling + + if self.output_cache: + return logits, output_k_cache, output_v_cache + return logits + + def get_example_inputs(self): + hidden_states = torch.randn( + (self.max_batch_size, self.ar_len, self.dim), dtype=torch.float32 + ) + inputs = list(super().get_example_inputs()) + inputs[0] = hidden_states + return tuple(inputs) + + def get_metadata(self): + meta_data = super().get_metadata() + if self.audio_token_id: + meta_data["audio_token_id"] = self.audio_token_id + if self.image_token_id: + meta_data["image_token_id"] = self.image_token_id + return meta_data + + +class MultiScopeAwareLlamaModel(LlamaModel): + def __init__( + self, + config: ModelArgs, + ar_len=1, + output_new_cache_only=True, + output_cache=True, + use_i64_token=False, + **kwargs, + ): + super().__init__( + config=config, + ar_len=ar_len, + output_new_cache_only=output_new_cache_only, + output_cache=output_cache, + use_i64_token=use_i64_token, + **kwargs, + ) + # Parameter final_logit_softcapping is not necessary for all + self.final_logit_softcapping = config.final_logit_softcapping + + # Gemma2/Gemma3 requires additional configuration parameters: + # - layer_types: Specifies the type of each layer (e.g., full vs. sliding attention) + # - local_rope_theta: Base frequency for local RoPE + # - sliding_window: Size of the sliding window for local attention + if config.layer_types is not None: + assert len(config.layer_types) == self.n_layers, ( + f"Length of layer_types ({len(config.layer_types)}) must match " + f"n_layers ({self.n_layers})" + ) + assert ( + config.local_rope_theta is not None + ), "local_rope_theta should not be None, please set it explicitly in config." + + self.sliding_window = config.sliding_window + + self.n_self_layers = config.n_layers - config.num_kv_shared_layers + if config.num_kv_shared_layers > 0: + self.embedding_scale_factor = math.sqrt(config.dim) + + # Per-Layer Embeddings (PLE) + self.use_per_layer_embedding = ( + config.vocab_size_per_layer_input and config.hidden_size_per_layer_input + ) + if self.use_per_layer_embedding: + self.embed_tokens_per_layer = nn.Embedding( + config.vocab_size_per_layer_input, + config.n_layers * config.hidden_size_per_layer_input, + ) + self.embed_scale_per_layer = math.sqrt(config.hidden_size_per_layer_input) + self.per_layer_model_projection = nn.Linear( + config.dim, + config.n_layers * config.hidden_size_per_layer_input, + bias=False, + ) + self.per_layer_projection_norm = torch.nn.RMSNorm( + config.hidden_size_per_layer_input, eps=config.norm_eps + ) + self._ple_input_scale = 1.0 / math.sqrt(2.0) + self._ple_proj_scale = 1.0 / math.sqrt(config.dim) + + def _compute_per_layer_inputs( + self, + tokens: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + n_layers = self.config.n_layers + hpl = self.config.hidden_size_per_layer_input + + per_layer_proj = ( + self.per_layer_model_projection(hidden_states) * self._ple_proj_scale + ) + per_layer_proj = per_layer_proj.view(bsz, seq_len, n_layers, hpl) + per_layer_proj = self.per_layer_projection_norm(per_layer_proj) + per_layer_embed = ( + self.embed_tokens_per_layer(tokens) * self.embed_scale_per_layer + ) + per_layer_embed = per_layer_embed.view(bsz, seq_len, n_layers, hpl) + + combined = (per_layer_proj + per_layer_embed) * self._ple_input_scale + return combined.permute(2, 0, 1, 3) + + def forward( # noqa: C901 + self, + tokens: torch.Tensor, + atten_mask: torch.Tensor, + window_atten_mask: torch.Tensor, + input_pos: Optional[torch.Tensor] = None, + *args, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + + output_k_cache = [] + output_v_cache = [] + # following tensors should be invariant across batches + global_freqs_cos = ( + self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos + ) + global_freqs_sin = ( + self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin + ) + local_freqs_cos = ( + self.local_freqs_cos[input_pos][0] + if self.use_kv_cache + else self.local_freqs_cos + ) + local_freqs_sin = ( + self.local_freqs_sin[input_pos][0] + if self.use_kv_cache + else self.local_freqs_sin + ) + + hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) + if self.use_per_layer_embedding: + per_layer_inputs = self._compute_per_layer_inputs(tokens, hidden_states) + + if self.config.num_kv_shared_layers > 0: + shared_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} + for ind, decoder_layer in enumerate(self.layers): + if self.config.is_sliding_attention(ind): + freqs_cos, freqs_sin, mask = ( + local_freqs_cos, + local_freqs_sin, + window_atten_mask, + ) + else: + freqs_cos, freqs_sin, mask = ( + global_freqs_cos, + global_freqs_sin, + atten_mask, + ) + per_layer_input = ( + per_layer_inputs[ind] if self.use_per_layer_embedding else None + ) + + if ind < self.n_self_layers: + # Self-decoder layer: own KV I/O. + k_caches = None + v_caches = None + if self.use_kv_cache: + offset_k = ind + offset_v = self.n_self_layers + offset_k + k_caches = args[offset_k] + v_caches = args[offset_v] + + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + atten_mask=mask, + k_caches=k_caches, + v_caches=v_caches, + per_layer_input=per_layer_input, + ) + output_k_cache.append(k) + output_v_cache.append(v) + + # Donor layers expose their full updated K/V for YOCO sharing. + if self.config.is_kv_donor_layer(ind): + if ( + decoder_layer.attention.output_new_cache_only + and k_caches is not None + ): + full_k = torch.cat([k_caches, k], dim=-1) + full_v = torch.cat([v_caches, v], dim=2) + else: + full_k, full_v = k, v + shared_kv[ind] = (full_k, full_v) + else: + # Cross-decoder layer: no KV I/O, consumes donor K/V. + kv_src = self.config.get_kv_shared_layer_index(ind) + donor_k, donor_v = None, None + if kv_src is not None and kv_src in shared_kv: + donor_k, donor_v = shared_kv[kv_src] + + hidden_states, _, _ = decoder_layer( + hidden_states, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + atten_mask=mask, + donor_k=donor_k, + donor_v=donor_v, + per_layer_input=per_layer_input, + ) + else: + for ind, decoder_layer in enumerate(self.layers): + k_caches = None + v_caches = None + if self.use_kv_cache: + offset_k = ind + offset_v = self.n_layers + offset_k + k_caches = args[offset_k] + v_caches = args[offset_v] + + if self.config.is_sliding_attention(ind): + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=local_freqs_cos, + freqs_sin=local_freqs_sin, + atten_mask=window_atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + else: + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=global_freqs_cos, + freqs_sin=global_freqs_sin, + atten_mask=atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + + output_k_cache.extend(k) + output_v_cache.extend(v) + + hidden_states = self.norm(hidden_states) + logits = self.output(hidden_states) + if self.final_logit_softcapping: + logits = logits / self.final_logit_softcapping + logits = torch.tanh(logits) + logits = logits * self.final_logit_softcapping + + if self.output_cache: + return logits, output_k_cache, output_v_cache + return logits + + def get_example_inputs(self): + inputs = list(super().get_example_inputs()) + causal_mask = CausalAttentionMask( + self.max_batch_size, self.ar_len, self.max_context_len + ) + sliding_window_mask = SlidingWindowAttentionMask( + self.max_batch_size, + self.ar_len, + self.max_context_len, + sliding_window=self.sliding_window, + ) + # Don't reverse the order of attention mask + inputs[1] = AttentionMask([causal_mask, sliding_window_mask]) + return tuple(inputs) + + def get_metadata(self): + meta_data = super().get_metadata() + meta_data["get_sliding_window"] = self.sliding_window + if self.config.global_head_dim: + meta_data["get_global_head_dim"] = self.config.global_head_dim + return meta_data + + def get_kv_head_dim(self): + return { + self.head_dim, + *filter(None, [self.config.global_head_dim]), + } diff --git a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp index 9b8cdd7999e..3556bf40da6 100644 --- a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp +++ b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp @@ -140,6 +140,18 @@ std::string get_formatted_prompt( formatted_prompt.append("\n"); formatted_prompt.append("model\n"); break; + case example::DecoderModelVersion::kGemma4: + formatted_prompt.append(""); + if (!system_prompt.empty()) { + formatted_prompt.append("<|turn>system\n"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("\n"); + } + formatted_prompt.append("<|turn>user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("\n"); + formatted_prompt.append("<|turn>model\n"); + break; case example::DecoderModelVersion::kGranite: if (!system_prompt.empty()) { formatted_prompt.append("<|start_of_role|>system<|end_of_role|>"); diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp index 6fb2255aa94..99f014fbabf 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp @@ -69,39 +69,66 @@ KVManager::KVManager(Metadata metadata, std::unique_ptr method_meta) Result attention_mask = method_meta->input_tensor_meta(1); attention_mask_dtype_ = attention_mask->scalar_type(); - // inputs are [input_tokens, attention_mask, (sliding window attention_mask), - // (input_pos), kv_caches] search kv_cache in inputs - for (int i = 2; i < method_meta->num_inputs(); i++) { - Result tensor_meta = method_meta->input_tensor_meta(i); - // k_cache: [1, n_heads, head_dim, seq_len] - size_t tensor_nbytes = tensor_meta->nbytes(); - size_t expected_tensor_nbytes = metadata_.head_dim * metadata_.num_heads * - metadata_.max_cache_len * getDtypeSize(tensor_meta->scalar_type()); - if (tensor_nbytes != expected_tensor_nbytes) { - // Not a kv_cache tensor (e.g. input_pos, sliding window attention mask). - continue; - } + // inputs: [tokens, atten_mask, (window_atten_mask), (input_pos), k_caches..., + // v_caches...]. + // outputs: [logits, k_cache_0..n-1, v_cache_0..n-1] + // k_cache shape: [bsz, n_kv_head, head_dim, ar_len] + // v_cache v_cache shape: [bsz, n_kv_head, ar_len, head_dim] + k_cache_.resize(metadata_.num_layers); + v_cache_.resize(metadata_.num_layers); + k_cache_in_nbytes_.resize(metadata_.num_layers); + k_cache_out_nbytes_.resize(metadata_.num_layers); + v_cache_in_nbytes_.resize(metadata_.num_layers); + v_cache_out_nbytes_.resize(metadata_.num_layers); + + for (int layer = 0; layer < metadata_.num_layers; ++layer) { + // k output: index 1 + layer + Result k_out = method_meta->output_tensor_meta(1 + layer); + // v output: index 1 + num_layers + layer + Result v_out = + method_meta->output_tensor_meta(1 + metadata_.num_layers + layer); + if (kv_cache_dtype_ == executorch::aten::ScalarType::Undefined) { - kv_cache_dtype_ = tensor_meta->scalar_type(); - } else { - ET_CHECK_MSG( - tensor_meta->scalar_type() == kv_cache_dtype_, - "Currently mixed scalar type of kv_cache is not allowed"); + kv_cache_dtype_ = k_out->scalar_type(); } + ET_CHECK_MSG( + k_out->scalar_type() == kv_cache_dtype_, + "Mixed scalar type of kv_cache is not allowed (k layer %d)", + layer); + ET_CHECK_MSG( + v_out->scalar_type() == kv_cache_dtype_, + "Mixed scalar type of kv_cache is not allowed (v layer %d)", + layer); + + // k_cache shape: [1, n_kv, head_dim, ar_len] -> head_dim at dim 2 + // v_cache shape: [1, n_kv, ar_len, head_dim] -> head_dim at dim 3 + const int64_t k_head_dim = k_out->sizes()[2]; + const int64_t v_head_dim = v_out->sizes()[3]; + k_head_dim_.push_back(k_head_dim); + v_head_dim_.push_back(v_head_dim); + + const size_t dtype_size = getDtypeSize(kv_cache_dtype_); + k_cache_out_nbytes_[layer] = + metadata_.num_heads * k_head_dim * metadata_.max_ar_len * dtype_size; + v_cache_out_nbytes_[layer] = + metadata_.num_heads * v_head_dim * metadata_.max_ar_len * dtype_size; + k_cache_in_nbytes_[layer] = + metadata_.num_heads * k_head_dim * metadata_.max_cache_len * dtype_size; + v_cache_in_nbytes_[layer] = + metadata_.num_heads * v_head_dim * metadata_.max_cache_len * dtype_size; } + ET_CHECK_MSG( kv_cache_dtype_ != executorch::aten::ScalarType::Undefined, - "kv_cache_dtype was not detected from method inputs"); - k_cache_.resize(metadata_.num_layers); - v_cache_.resize(metadata_.num_layers); + "kv_cache_dtype was not detected from method outputs"); - // Calculate cache size - size_t cache_in_bytes = metadata_.num_layers * metadata_.num_heads * - metadata_.head_dim * metadata_.max_cache_len * - getDtypeSize(kv_cache_dtype_); - size_t cache_out_bytes = metadata_.num_layers * metadata_.num_heads * - metadata_.head_dim * metadata_.max_ar_len * getDtypeSize(kv_cache_dtype_); - total_cache_size_ = 2 * (cache_in_bytes + cache_out_bytes); + // Total cache size across all layers (K + V, input + output) + total_cache_size_ = 0; + for (int layer = 0; layer < metadata_.num_layers; ++layer) { + total_cache_size_ += k_cache_in_nbytes_[layer] + + k_cache_out_nbytes_[layer] + v_cache_in_nbytes_[layer] + + v_cache_out_nbytes_[layer]; + } }; void KVManager::init_attention_mask( @@ -269,15 +296,15 @@ void KVManager::update_attention_mask( void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { cur_ar_len_ = ar_len; - const size_t cache_in_bytes = metadata_.num_heads * metadata_.head_dim * - metadata_.max_cache_len * getDtypeSize(kv_cache_dtype_); - const size_t cache_out_bytes = metadata_.num_heads * metadata_.head_dim * - metadata_.max_ar_len * getDtypeSize(kv_cache_dtype_); for (int layer = 0; layer < metadata_.num_layers; ++layer) { - k_cache_[layer].buffer = buffer_manager->allocate(cache_in_bytes); - k_cache_[layer].output_buffer = buffer_manager->allocate(cache_out_bytes); - v_cache_[layer].buffer = buffer_manager->allocate(cache_in_bytes); - v_cache_[layer].output_buffer = buffer_manager->allocate(cache_out_bytes); + k_cache_[layer].buffer = + buffer_manager->allocate(k_cache_in_nbytes_[layer]); + k_cache_[layer].output_buffer = + buffer_manager->allocate(k_cache_out_nbytes_[layer]); + v_cache_[layer].buffer = + buffer_manager->allocate(v_cache_in_nbytes_[layer]); + v_cache_[layer].output_buffer = + buffer_manager->allocate(v_cache_out_nbytes_[layer]); } } @@ -286,14 +313,17 @@ void KVManager::rearrange_cache(int32_t ar_len_dst) { if (cur_ar_len_ == ar_len_dst) return; for (int layer = 0; layer < metadata_.num_layers; ++layer) { - rearrange_key(k_cache_[layer], ar_len_dst); - rearrange_value(v_cache_[layer], ar_len_dst); + rearrange_key(k_cache_[layer], ar_len_dst, k_head_dim_[layer]); + rearrange_value(v_cache_[layer], ar_len_dst, v_head_dim_[layer]); } // rearrange done. cur_ar_len_ = ar_len_dst; } -void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { +void KVManager::rearrange_key( + KVCache& k_cache, + int32_t ar_len_dst, + int64_t head_dim) { const int32_t src_cache_num = (cur_ar_len_ == metadata_.context_len) ? metadata_.context_len : metadata_.context_len - cur_ar_len_; @@ -304,18 +334,18 @@ void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { size_t dst_cache_nbytes = dst_cache_num * getDtypeSize(kv_cache_dtype_); if (src_cache_num > dst_cache_num) { // copy from first dimension - for (int i = 0; i < metadata_.head_dim * metadata_.num_heads; i++) { + for (int i = 0; i < head_dim * metadata_.num_heads; i++) { std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_nbytes); k_cache_in_read_ptr += src_cache_nbytes; k_cache_in_write_ptr += dst_cache_nbytes; } } else { k_cache_in_read_ptr += - (metadata_.head_dim * metadata_.num_heads - 1) * src_cache_nbytes; + (head_dim * metadata_.num_heads - 1) * src_cache_nbytes; k_cache_in_write_ptr += - (metadata_.head_dim * metadata_.num_heads - 1) * dst_cache_nbytes; + (head_dim * metadata_.num_heads - 1) * dst_cache_nbytes; // copy from last dimension - for (int i = 0; i < metadata_.head_dim * metadata_.num_heads; i++) { + for (int i = 0; i < head_dim * metadata_.num_heads; i++) { std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_nbytes); k_cache_in_read_ptr -= src_cache_nbytes; k_cache_in_write_ptr -= dst_cache_nbytes; @@ -323,7 +353,10 @@ void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { } } -void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) { +void KVManager::rearrange_value( + KVCache& v_cache, + int32_t ar_len_dst, + int64_t head_dim) { const int32_t src_cache_num = (cur_ar_len_ == metadata_.context_len) ? metadata_.context_len : metadata_.context_len - cur_ar_len_; @@ -338,23 +371,23 @@ void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) { std::memmove( v_cache_in_write_ptr, v_cache_in_read_ptr, - dst_cache_nbytes * metadata_.head_dim); - v_cache_in_read_ptr += src_cache_nbytes * metadata_.head_dim; - v_cache_in_write_ptr += dst_cache_nbytes * metadata_.head_dim; + dst_cache_nbytes * head_dim); + v_cache_in_read_ptr += src_cache_nbytes * head_dim; + v_cache_in_write_ptr += dst_cache_nbytes * head_dim; } } else { v_cache_in_read_ptr += - metadata_.head_dim * (metadata_.num_heads - 1) * src_cache_nbytes; + head_dim * (metadata_.num_heads - 1) * src_cache_nbytes; v_cache_in_write_ptr += - metadata_.head_dim * (metadata_.num_heads - 1) * dst_cache_nbytes; + head_dim * (metadata_.num_heads - 1) * dst_cache_nbytes; // copy from last dimension for (int i = 0; i < metadata_.num_heads; i++) { std::memmove( v_cache_in_write_ptr, v_cache_in_read_ptr, - src_cache_nbytes * metadata_.head_dim); - v_cache_in_read_ptr -= src_cache_nbytes * metadata_.head_dim; - v_cache_in_write_ptr -= dst_cache_nbytes * metadata_.head_dim; + src_cache_nbytes * head_dim); + v_cache_in_read_ptr -= src_cache_nbytes * head_dim; + v_cache_in_write_ptr -= dst_cache_nbytes * head_dim; } } } @@ -370,8 +403,9 @@ void KVManager::update_cache( cur_ar_len_, ar_len); for (int layer = 0; layer < metadata_.num_layers; ++layer) { - update_key(k_cache_[layer], n_past, n_update, selected); - update_value(v_cache_[layer], n_past, n_update, selected); + update_key(k_cache_[layer], n_past, n_update, selected, k_head_dim_[layer]); + update_value( + v_cache_[layer], n_past, n_update, selected, v_head_dim_[layer]); } } @@ -379,7 +413,8 @@ void KVManager::update_key( KVCache& k_cache, int32_t n_past, int32_t n_update, - const std::vector& selected) { + const std::vector& selected, + int64_t head_dim) { std::byte* write_ptr = k_cache.buffer; std::byte* read_ptr = k_cache.output_buffer; const int32_t copy_size = n_update * getDtypeSize(kv_cache_dtype_); @@ -388,7 +423,7 @@ void KVManager::update_key( : (metadata_.context_len - cur_ar_len_) * getDtypeSize(kv_cache_dtype_); const int32_t out_size = cur_ar_len_ * getDtypeSize(kv_cache_dtype_); const int32_t past_size = n_past * getDtypeSize(kv_cache_dtype_); - const int32_t n_iter = metadata_.head_dim * metadata_.num_heads; + const int32_t n_iter = head_dim * metadata_.num_heads; write_ptr += past_size; if (selected.empty()) { @@ -423,21 +458,19 @@ void KVManager::update_value( KVCache& v_cache, int32_t n_past, int32_t n_update, - const std::vector& selected) { + const std::vector& selected, + int64_t head_dim) { std::byte* write_ptr = v_cache.buffer; std::byte* read_ptr = v_cache.output_buffer; - const int32_t copy_size = - n_update * metadata_.head_dim * getDtypeSize(kv_cache_dtype_); - const int32_t past_size = - n_past * metadata_.head_dim * getDtypeSize(kv_cache_dtype_); + const int32_t copy_size = n_update * head_dim * getDtypeSize(kv_cache_dtype_); + const int32_t past_size = n_past * head_dim * getDtypeSize(kv_cache_dtype_); const int32_t n_iter = metadata_.num_heads; const int32_t iter_size = (cur_ar_len_ == metadata_.context_len) - ? metadata_.context_len * metadata_.head_dim * - getDtypeSize(kv_cache_dtype_) - : (metadata_.context_len - cur_ar_len_) * metadata_.head_dim * + ? metadata_.context_len * head_dim * getDtypeSize(kv_cache_dtype_) + : (metadata_.context_len - cur_ar_len_) * head_dim * getDtypeSize(kv_cache_dtype_); const int32_t out_size = - cur_ar_len_ * metadata_.head_dim * getDtypeSize(kv_cache_dtype_); + cur_ar_len_ * head_dim * getDtypeSize(kv_cache_dtype_); write_ptr += past_size; @@ -453,14 +486,13 @@ void KVManager::update_value( auto wp = write_ptr, rp = read_ptr; for (auto sel : selected) { if (sel) { - std::memcpy( - wp, rp, metadata_.head_dim * getDtypeSize(kv_cache_dtype_)); - wp += metadata_.head_dim * getDtypeSize(kv_cache_dtype_); + std::memcpy(wp, rp, head_dim * getDtypeSize(kv_cache_dtype_)); + wp += head_dim * getDtypeSize(kv_cache_dtype_); update_times--; if (update_times == 0) break; } - rp += metadata_.head_dim * getDtypeSize(kv_cache_dtype_); + rp += head_dim * getDtypeSize(kv_cache_dtype_); } write_ptr += iter_size; read_ptr += out_size; diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h index 3b8e67dd38d..cbebaa5d8cb 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h +++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h @@ -171,21 +171,23 @@ class KVManager { private: // Helper functions to rearrange and update key and value caches - void rearrange_key(KVCache& k_cache, int32_t ar_len_dst); + void rearrange_key(KVCache& k_cache, int32_t ar_len_dst, int64_t head_dim); - void rearrange_value(KVCache& v_cache, int32_t ar_len_dst); + void rearrange_value(KVCache& v_cache, int32_t ar_len_dst, int64_t head_dim); void update_key( KVCache& k_cache, int32_t n_past, int32_t n_update, - const std::vector& selected); + const std::vector& selected, + int64_t head_dim); void update_value( KVCache& v_cache, int32_t n_past, int32_t n_update, - const std::vector& selected); + const std::vector& selected, + int64_t head_dim); // metadata Metadata metadata_; @@ -200,5 +202,13 @@ class KVManager { // output: layer -> head * head_dim * max_ar_len std::vector k_cache_; std::vector v_cache_; + // Per-layer buffer sizes and head_dim (may differ when head_dim varies across + // layers) + std::vector k_cache_in_nbytes_; + std::vector k_cache_out_nbytes_; + std::vector v_cache_in_nbytes_; + std::vector v_cache_out_nbytes_; + std::vector k_head_dim_; + std::vector v_head_dim_; }; } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp index 0e1a27a751a..409f7c96f34 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp @@ -126,6 +126,9 @@ Runner::Runner( } else if (decoder_model_version == "gemma3") { decoder_model_version_ = DecoderModelVersion::kGemma3; cache_mode_ = CacheMode::HybridCache; + } else if (decoder_model_version == "gemma4") { + decoder_model_version_ = DecoderModelVersion::kGemma4; + cache_mode_ = CacheMode::HybridCache; } else if (decoder_model_version == "granite") { decoder_model_version_ = DecoderModelVersion::kGranite; } else if (decoder_model_version == "phi_4_mini") { @@ -208,6 +211,8 @@ Error Runner::load() { decoder_model_version_ == DecoderModelVersion::kGemma2 || decoder_model_version_ == DecoderModelVersion::kGemma3) { eos_ids->insert(tokenizer_->encode("", 0, 0).get()[0]); + } else if (decoder_model_version_ == DecoderModelVersion::kGemma4) { + eos_ids->insert(tokenizer_->encode("", 0, 0).get()[0]); } else if (decoder_model_version_ == DecoderModelVersion::kCodegen) { eos_ids->insert(tokenizer_->encode("<|endoftext|>", 0, 0).get()[0]); } else if (decoder_model_version_ == DecoderModelVersion::kGlm) { @@ -274,6 +279,14 @@ Error Runner::load() { sliding_window_evalue__, module_->get("get_sliding_window")); sliding_window = sliding_window_evalue__.toInt(); } + // Gemma4 uses YOCO: only self-decoder layers have KV I/O; use + // get_n_self_layers when available so KVManager allocates the correct number + // of caches. + int64_t num_kv_layers = num_layers; + if (module_->method_names()->count("get_n_self_layers") > 0) { + ET_ASSIGN_OR_RETURN(num_layers_evalue__, module_->get("get_n_self_layers")); + num_kv_layers = num_layers_evalue__.toScalar().to(); + } kv_manager_ = std::make_unique( KVManager::Metadata{ context_len_, @@ -281,7 +294,7 @@ Error Runner::load() { max_ar_len, max_cache_len, num_heads, - num_layers}, + num_kv_layers}, std::make_unique( std::move(module_->method_meta(token_generator_method_name).get()))); @@ -299,7 +312,7 @@ Error Runner::load() { PromptProcessor::Metadata{ context_len_, num_heads, - num_layers, + num_kv_layers, prompt_processor_ar_len, vocab_size, use_int64_token, @@ -317,7 +330,7 @@ Error Runner::load() { LhdTokenGenerator::Metadata{ context_len_, num_heads, - num_layers, + num_kv_layers, token_generator_ar_len, vocab_size, use_int64_token, @@ -339,7 +352,7 @@ Error Runner::load() { TokenGenerator::Metadata{ context_len_, num_heads, - num_layers, + num_kv_layers, token_generator_ar_len, vocab_size, use_int64_token, diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.h b/examples/qualcomm/oss_scripts/llama/runner/runner.h index 86934558656..a996c92997d 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/runner.h +++ b/examples/qualcomm/oss_scripts/llama/runner/runner.h @@ -44,6 +44,7 @@ enum DecoderModelVersion { kCodegen, kGlm, kGemma2, + kGemma4, }; class Runner : public executorch::extension::llm::IRunner { diff --git a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py index 4eecf308f88..9f7188de807 100644 --- a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py +++ b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py @@ -335,6 +335,30 @@ def __init__(self, verbose: bool = False): self.recipe.custom_quant_annotations.append(annotate_kv_8bit) +class Gemma4QuantRecipe(StaticLLMQuantRecipe): + default_quant_dtype = QuantDtype.use_16a8w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ).add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 64, 1, 1)}, + ) + + class GLM_1_5B_InstructQuantRecipe(StaticLLMQuantRecipe): default_quant_dtype = QuantDtype.use_16a4w diff --git a/examples/qualcomm/oss_scripts/llama/tokenizer.py b/examples/qualcomm/oss_scripts/llama/tokenizer.py index 954b73384fa..7fab6a13ccc 100644 --- a/examples/qualcomm/oss_scripts/llama/tokenizer.py +++ b/examples/qualcomm/oss_scripts/llama/tokenizer.py @@ -15,7 +15,6 @@ AUDIO_ENCODER, VISION_ENCODER, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs from pytorch_tokenizers import get_tokenizer, TiktokenTokenizer from pytorch_tokenizers.llama2c import Llama2cTokenizer as SentencePieceTokenizer @@ -77,8 +76,9 @@ def __init__(self, control_args: argparse.Namespace, config: LLMModelConfig): config.params_path if control_args.params is None else control_args.params ) with open(params_path) as f: - model_args = ModelArgs(**json.load(f)) - self.vocab_size = model_args.vocab_size + model_args = json.load(f) + + self.vocab_size = model_args.get("text_config", model_args)["vocab_size"] self.runtime_tokenizer_path = self._init_tokenizer( control_args.tokenizer_model, control_args.tokenizer_bin @@ -136,6 +136,10 @@ def _from_hf(self): # Override the default BOS and EOS token IDs for codegen2_1b tokenizer.bos_id = 1 tokenizer.eos_id = 2 + if self.decoder_model == "gemma4-e2b": + # Override the default BOS and EOS token IDs for gemma4_e2b + tokenizer.bos_id = 2 + tokenizer.eos_id = 1 return runtime_tokenizer_path, tokenizer, chat_template diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py index da3a165277f..afd34eeed57 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/attention_sink_wrappers.py @@ -43,7 +43,7 @@ ATTENTION_SINK_EVICTOR, DECODER_GRAPH_NAMES, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( +from executorch.examples.qualcomm.oss_scripts.llama.model import ( AttentionSinkRope, ModelArgs, ) diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py b/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py index 504a3d3e9e4..787232f58da 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/base_component.py @@ -28,7 +28,7 @@ AUDIO_ENCODER, VISION_ENCODER, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs +from executorch.examples.qualcomm.oss_scripts.llama.model import ModelArgs from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import ( StaticLLMQuantRecipe, ) diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py index 9a1ad50f854..4d3715b0a26 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py @@ -78,12 +78,10 @@ PerLayerSqnrAnalyzer, save_suggest_recipes, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.embedding import ( - TokenEmbedding, -) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( +from executorch.examples.qualcomm.oss_scripts.llama.model import ( LlamaModel, ModelArgs, + TokenEmbedding, ) from executorch.examples.qualcomm.oss_scripts.llama.quantize import PTQStrategy from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import ( @@ -187,9 +185,22 @@ def __init__( params_path = ( config.params_path if control_args.params is None else control_args.params ) + if control_args.decoder_model == "gemma4-e2b": + from executorch.examples.qualcomm.oss_scripts.llama import ( + load_gemma4_model_args, + ) + + model_args = load_gemma4_model_args(params_path) + else: + with open(params_path) as f: + model_args = ModelArgs(**json.load(f)) with open(params_path) as f: self.model_args = process_model_args( - control_args, ModelArgs(**json.load(f)), self.quant_recipe, config, mode + control_args, + model_args, + self.quant_recipe, + config, + mode, ) # prepare instance self.tok_embedding, self.decoder = self._prepare_model() @@ -308,8 +319,8 @@ def permute(w, heads, partial_rotary_dim): for layer in decoder.layers: if getattr(layer.attention, "prepare_attention_conv", None): layer.attention.prepare_attention_conv() - if getattr(layer.feed_forward, "prepare_feedfoward_conv", None): - layer.feed_forward.prepare_feedfoward_conv() + if getattr(layer.feed_forward, "prepare_feedforward_conv", None): + layer.feed_forward.prepare_feedforward_conv() decoder = convert_linear_to_conv2d(decoder) @@ -320,16 +331,10 @@ def permute(w, heads, partial_rotary_dim): # check embedding fallback if self.control_args.embedding_quantize: - decoder = get_quant_embedding_transform( - embedding_quantize=self.control_args.embedding_quantize - )(decoder) self.passes_job[I64toI32][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ "skip_node" ] = {"tokens"} if self.apply_embedding: - tok_embedding = get_quant_embedding_transform( - embedding_quantize=self.control_args.embedding_quantize - )(tok_embedding) self.tok_embedding_passes_job[I64toI32][ QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY ]["skip_node"] = {"tokens"} @@ -398,15 +403,8 @@ def _get_model_instance(self) -> LlamaModel: decoder.vocab_size, ), } - # shape of k caches and v caches - self.kv_cache_shape = { - # single head, kv input - (self.meta["get_head_dim"], self.meta["get_max_context_len"]), - (self.meta["get_max_context_len"], self.meta["get_head_dim"]), - # single head, kv output - (self.meta["get_head_dim"], self.meta["get_ar_len"]), - (self.meta["get_ar_len"], self.meta["get_head_dim"]), - } + self.kv_cache_shape = decoder.get_kv_cache_shapes() + self.freq_shape = decoder.get_freq_shapes() if self.apply_embedding: self.tok_embedding_export_input = ( @@ -462,10 +460,6 @@ def _tag_ios(self, node, fixed_point_type): ), } - freq_shape = { - (self.meta["get_ar_len"], self.meta["get_head_dim"] // 2), - } - freq_op = { exir_ops.edge.aten.select.int, } @@ -473,7 +467,8 @@ def _tag_ios(self, node, fixed_point_type): if node.op == "placeholder": if ( - len(users := list(node.users)) == 1 + len(users := list(node.users)) > 0 + and "args_" in node.name and users[0].meta["val"].size()[-2:] in self.kv_cache_shape ): quant_io_type = fixed_point_type["kv_type"] @@ -495,7 +490,22 @@ def _tag_ios(self, node, fixed_point_type): quant_io_type = fixed_point_type["io_type"] # tag select op as quantized tensors for freq_sin and freq_cos. It is caused by sharding - if node.target in freq_op and node.meta["val"].size() in freq_shape: + if node.target in freq_op and node.meta["val"].size() in self.freq_shape: + quant_io_type = fixed_point_type["io_type"] + + # tag per-layer embedding. + if ( + self.control_args.decoder_model == "gemma4-e2b" + and "stack_trace" in node.meta + and "per_layer_inputs[ind]" in node.meta["stack_trace"] + ): + quant_io_type = fixed_point_type["io_type"] + + # tag donor kv + if self.control_args.decoder_model == "gemma4-e2b" and ( + is_node_src_start_with_name(node, "full_k") + or is_node_src_start_with_name(node, "full_v") + ): quant_io_type = fixed_point_type["io_type"] return quant_io_type @@ -566,6 +576,16 @@ def quantize(self, request: Request): # noqa: C901 f"unknown logits io bit width {self.quant_recipe.get_logits_output_bit_width()}" ) + # embedding fallback and quantization + if self.control_args.embedding_quantize: + self.decoder = get_quant_embedding_transform( + embedding_quantize=self.control_args.embedding_quantize + )(self.decoder) + if self.apply_embedding: + self.tok_embedding = get_quant_embedding_transform( + embedding_quantize=self.control_args.embedding_quantize + )(self.tok_embedding) + data = request.method_data[TEXT_DECODER] quantizer = make_quantizer(backend=data.backend, soc_model=data.soc_model) @@ -852,7 +872,12 @@ def _set_attr( if "args_" in node.name: args_idx = int(node.name.split("_")[-1]) - if args_idx >= self.decode.meta["get_n_layers"]: + n_cache_layers = ( + self.decode.meta["get_n_self_layers"] + if self.control_args.decoder_model == "gemma4-e2b" + else self.decode.meta["get_n_layers"] + ) + if args_idx >= n_cache_layers: v_input_cache_nodes.append(node) else: k_input_cache_nodes.append(node) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index 75097718032..fcd6a9504db 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -884,10 +884,12 @@ def create_submodule_from_nodes( # all uses with a getitem call to the 0th index of the result with gm.graph.inserting_after(submodule_node): proxy_out = torch.fx.Proxy(submodule_node)[0].node # type: ignore[index] - submodule_node.replace_all_uses_with(proxy_out) + submodule_node.replace_all_uses_with(proxy_out, propagate_meta=True) proxy_out.meta["val"] = submodule_node.meta["val"] # Reset the args since it was overwritten in the previous line proxy_out.args = (submodule_node, 0) + proxy_out.meta.pop("nn_module_stack", None) + proxy_out.meta.pop("source_fn_stack", None) else: # fuse_as_graphmodule will automatically propagate the metadata of the # partition's last node to the getitem nodes that appear after the