Views
No views yet
wget https://raw.githubusercontent.com/JT-Ushio/MHA2MLA/refs/heads/main/src/mha2mla/monkey_patch.pyqk_tensor_360M.pth as an example:wget https://github.com/JT-Ushio/MHA2MLA/raw/refs/heads/main/utils/qk_tensor_360M.pthfnlp/SmolLM-360M-GQA-d_kv_128 as an example:1import torch
2from transformers import AutoConfig, AutoTokenizer, LlamaForCausalLM
3from monkey_patch import infer_monkey_patch
4
5model_name = "fnlp/SmolLM-360M-GQA-d_kv_128"
6
7# Monkey Patch: MHA -> MLA
8config = AutoConfig.from_pretrained(model_name)
9if "RoPE" in config:
10 config.RoPE["qk_tensor_path"] = "qk_tensor_360M.pth" # Configuration for Specific Models
11 infer_monkey_patch(config.RoPE)
12
13tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
14model = LlamaForCausalLM.from_pretrained(model_name, config=config, torch_dtype=torch.bfloat16).cuda()
15
16# Generate
17text = "Which American-born Sinclair won the Nobel Prize for Literature in 1930?"
18inputs = tokenizer(text, return_tensors="pt").to(model.device)
19generation_kwargs = {"do_sample": False, "use_cache": True, "max_new_tokens": 128}
20output = model.generate(**inputs, **generation_kwargs)
21
22print(tokenizer.decode(output[0], skip_special_tokens=True))
23# - Sinclair Lewis@misc{ji2025economicalinferenceenablingdeepseeks,
title={Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs},
author={Tao Ji and Bin Guo and Yuanbin Wu and Qipeng Guo and Lixing Shen and Zhan Chen and Xipeng Qiu and Qi Zhang and Tao Gui},
year={2025},
eprint={2502.14837},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2502.14837},
}