chore: initialize the app

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

@@ -0,0 +1 @@
+~

README.md

@@ -0,0 +1,45 @@
+---
+title: Ghost 8B Beta (128k)
+emoji: 👻 / 📚
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: 4.36.1
+app_file: app.py
+pinned: false
+suggested_hardware: a10g-small
+language:
+  - en
+  - vi
+  - es
+  - pt
+  - de
+  - it
+  - fr
+  - ko
+  - zh
+license: other
+license_name: ghost-llms
+license_link: https://ghost-x.org/ghost-llms-license
+tags:
+  - ghost
+---
+
+# ~
+
+### Notes
+
+The extension source code belongs to: "LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning".
+
+See source code details [here](https://github.com/datamllab/LongLM).
+
+```
+@misc{jin2024llm,
+      title={LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning}, 
+      author={Hongye Jin and Xiaotian Han and Jingfeng Yang and Zhimeng Jiang and Zirui Liu and Chia-Yuan Chang and Huiyuan Chen and Xia Hu},
+      year={2024},
+      eprint={2401.01325},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```

SelfExtend.py

@@ -0,0 +1,199 @@
+from types import MethodType
+from functools import partial
+import self_extend_patch as SE
+
+def modify_method_of_instance(instance, target_class_name, target_method_name, new_method, visited_instances=None):
+    """
+        This function modifies the method of an instance of a model class. 
+        It's part from chat-GPT.
+        It will replace the method  with the new method.
+        Currently, we only use this function to modify the attention method of a model. Do not test it further. 
+
+        instance: 
+            instance of a model to modify.
+        target_class_name: 
+            name of the attention class to modify. E.g. 'LlamaAttention', 'GPTNeoXAttention', etc.
+        new_method: new method to replace the original method. E.g. 'self_extend_forward'. 
+            It should include a parameter 'self' to be binded to the instance.
+    """
+    target_found = False
+    if visited_instances is None:
+        visited_instances = set()
+    # Unique identifier for the instance (using id() since object's id is unique)
+    instance_id = id(instance)
+    if instance_id in visited_instances:
+        target_found = False
+        return target_found
+    # Add the instance to the already_visited set
+    visited_instances.add(instance_id)
+
+    # Check if this instance is of the target class
+    if instance.__class__.__name__ == target_class_name:
+        bond_method = MethodType(new_method, instance) 
+        setattr(instance, target_method_name, bond_method)
+        target_found = True
+        return target_found
+    elif hasattr(instance, '__dict__'):
+        for attr_name, attr_value in instance.__dict__.items():
+            if isinstance(attr_value, object) and not isinstance(attr_value, (list, tuple, dict, set)):
+                _found = modify_method_of_instance(attr_value, target_class_name, target_method_name, new_method, visited_instances)
+                if _found:
+                    target_found = True
+            elif isinstance(attr_value, (list, tuple)):
+                for item in attr_value:
+                    if isinstance(item, object):
+                        _found = modify_method_of_instance(item, target_class_name, target_method_name, new_method, visited_instances)
+                        if _found:
+                            target_found = True
+            # If attribute value is a dictionary, iterate over its values and recurse
+            # E.g, for a ModuleList, its moudels are stored in a dictionary: ._modules
+            elif isinstance(attr_value, dict):
+                for key, value in attr_value.items():
+                    if isinstance(value, object):
+                        _found = modify_method_of_instance(value, target_class_name, target_method_name, new_method, visited_instances)
+                        if _found:
+                            target_found = True
+            # If attribute value is a set, iterate and recurse
+            elif isinstance(attr_value, set):
+                for item in attr_value:
+                    if isinstance(item, object):
+                        _found = modify_method_of_instance(item, target_class_name, target_method_name, new_method, visited_instances)
+                        if _found:
+                            target_found = True
+
+    return target_found
+
+
+def apply(loaded_model, group_size, window_size, enable_flash_attention=False, scale_base=-1, flash_attention_impl="triton"):
+    '''
+        loaded_model: 
+            model to apply the self-attention extension. 
+        group_size: 
+            group size for the self-attention extension. 
+        window_size: 
+            window size for the self-attention extension. 
+        scale_base:
+            base for the scale, equal to pretraining length. 
+            e.g. 4096 for Llama, 8192 for Gemma
+
+            Two recommended scale factor:
+                yarn: https://arxiv.org/abs/2309.00071
+                log: https://arxiv.org/abs/2202.12172 ; https://kexue.fm/archives/8823
+            This is helpful while retrieving a long sequence (e.g a long passkey).
+            But on real-world data, the impact is minor. (e.g. on LongBench, LEval).
+
+            The reported results in our paper does not use this scale except for long passkey retrieval.
+    '''
+    arch_name = loaded_model.__class__.__name__
+    if 'Llama' in arch_name:
+        if enable_flash_attention:
+            if flash_attention_impl == "flash_attn":
+                self_extend_attention_forward = partial(SE.Llama.flash_self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+                modifed_1 = modify_method_of_instance(loaded_model, "LlamaFlashAttention2", "_flash_attention_forward", SE.selfextend_flash_attn.flash_attention2_forward_with_window_size)
+                modifed_2 = modify_method_of_instance(loaded_model, "LlamaFlashAttention2", "forward", self_extend_attention_forward)
+                print("Using flash_attn flash self_extend!!")
+                if (not modifed_1) or (not modifed_2):
+                    raise Exception(f"Failed to modify the attention method of {arch_name}")
+
+            elif flash_attention_impl == "triton":
+                self_extend_attention_forward = partial(SE.Llama.flash_self_extend_forward_triton,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+                modifed = modify_method_of_instance(loaded_model, "LlamaFlashAttention2", "forward", self_extend_attention_forward)
+                print("Using triton flash self_extend!!")
+                if (not modifed):
+                    raise Exception(f"Failed to modify the attention method of {arch_name}")
+            else:
+                raise Exception(f"Need to set the flash_attention_impl to 'flash_attn' or 'triton'.")
+
+
+        else:
+            self_extend_attention_forward = partial(SE.Llama.self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            # after the default version of attention in 4.36 is LlamaSpdaAttention, but in before 4,36 or in 4.38, it is LlamaAttention
+            # print("loaded_model", loaded_model)
+            modifed_2 = modify_method_of_instance(loaded_model, "LlamaAttention", "forward", self_extend_attention_forward)
+            if not modifed_2:
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+    elif 'Mistral' in arch_name:
+        # Mistral shares the same architecture with Llama, so the implementation should be exchangable.
+        if enable_flash_attention:
+            self_extend_attention_forward = partial(SE.Mistral.flash_self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_1 = modify_method_of_instance(loaded_model, "MistralFlashAttention2", "_flash_attention_forward", SE.selfextend_flash_attn.flash_attention2_forward_with_window_size)
+            modifed_2 = modify_method_of_instance(loaded_model, "MistralFlashAttention2", "forward", self_extend_attention_forward)
+            if (not modifed_1) or (not modifed_2):
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+        else:
+            self_extend_attention_forward = partial(SE.Mistral.self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_2 = modify_method_of_instance(loaded_model, "MistralAttention", "forward", self_extend_attention_forward)
+            if not modifed_2:
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+    elif 'Gemma' in arch_name:
+        if enable_flash_attention:
+            self_extend_attention_forward = partial(SE.Gemma.flash_self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_1 = modify_method_of_instance(loaded_model, "GemmaFlashAttention2", "_flash_attention_forward", SE.selfextend_flash_attn.flash_attention2_forward_with_window_size)
+            modifed_2 = modify_method_of_instance(loaded_model, "GemmaFlashAttention2", "forward", self_extend_attention_forward)
+            if (not modifed_1) or (not modifed_2):
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+        else:
+            self_extend_attention_forward = partial(SE.Gemma.self_extend_forward,
+                                            group_size_1=group_size,
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_2= modify_method_of_instance(loaded_model, "GemmaAttention", "forward", self_extend_attention_forward)
+            if not modifed_2:
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+    elif 'Qwen2' in arch_name:
+        if enable_flash_attention:
+            self_extend_attention_forward = partial(SE.Qwen2.flash_self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_1 = modify_method_of_instance(loaded_model, "Qwen2FlashAttention2", "_flash_attention_forward", SE.selfextend_flash_attn.flash_attention2_forward_with_window_size)
+            modifed_2 = modify_method_of_instance(loaded_model, "Qwen2FlashAttention2", "forward", self_extend_attention_forward)
+            if (not modifed_1) or (not modifed_2):
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+        else:
+            self_extend_attention_forward = partial(SE.Qwen2.self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_2 = modify_method_of_instance(loaded_model, "Qwen2Attention", "forward", self_extend_attention_forward)
+            if not modifed_2:
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+    elif 'Phi' in arch_name:
+        if enable_flash_attention:
+            self_extend_attention_forward = partial(SE.Phi.flash_self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_1 = modify_method_of_instance(loaded_model, "PhiFlashAttention2", "_flash_attention_forward", SE.selfextend_flash_attn.flash_attention2_forward_with_window_size)
+            modifed_2 = modify_method_of_instance(loaded_model, "PhiFlashAttention2", "forward", self_extend_attention_forward)
+            if (not modifed_1) or (not modifed_2):
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+        else:
+            self_extend_attention_forward = partial(SE.Phi.self_extend_forward,
+                                            group_size_1=group_size, 
+                                            group_size_2=window_size,
+                                            scale_base=scale_base)
+            modifed_2 = modify_method_of_instance(loaded_model, "PhiAttention", "forward", self_extend_attention_forward)
+            if not modifed_2:
+                raise Exception(f"Failed to modify the attention method of {arch_name}")
+    else:
+        raise NotImplementedError
+

app.py

@@ -0,0 +1,388 @@
+# pylint: skip-file
+
+import subprocess
+
+subprocess.run(
+    f"pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
+import os
+from threading import Thread
+from typing import Iterator
+
+import gradio as gr
+import spaces
+import torch
+import SelfExtend
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+
+MAX_MAX_NEW_TOKENS = 4096
+DEFAULT_MAX_NEW_TOKENS = 1536
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "123392"))
+
+DESCRIPTION = """\
+# Playground with Ghost 8B Beta (p)
+
+**Ghost 8B Beta** is a large language model developed with goals that include excellent multilingual support, superior knowledge capabilities, and cost-effectiveness. The model comes in two context length versions, 8k and 128k, along with multilingual function tools support by default. 
+
+The languages supported are 🇺🇸 English, 🇫🇷 French, 🇮🇹 Italian, 🇪🇸 Spanish, 🇵🇹 Portuguese, 🇩🇪 German, 🇻🇳 Vietnamese, 🇰🇷 Korean and 🇨🇳 Chinese.
+
+📋 Note: current model version is "disl-0x5" (10 Jul 2024), context length 128k (123392 tokens) and current status is "moderating / previewing". For detailed information about the model, see [here](https://ghost-x.org/docs/models/ghost-8b-beta/). Try to experience it the way you want!
+"""
+
+
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+   <h1 style="font-size: 26px; margin-bottom: 2px; opacity: 0.20;">👻 Ghost 8B Beta</h1>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.10;">Ask and share whatever you want ~</p>
+</div>
+"""
+
+LICENSE = """
+<p/>
+
+---
+Ghost 8B Beta may give inaccurate information, including information about people, so please verify Ghost 8B Beta's answers. [Ghost 8B Beta](https://ghost-x.org/docs/models/ghost-8b-beta/) by [Ghost X](https://ghost-x.org).
+"""
+
+EXAMPLES = [
+    [
+        "What is the significance of the Higgs boson in the Standard Model of particle physics?"
+    ],
+    [
+        "Qu'est-ce que l'effet fondateur et comment influence-t-il la diversité génétique d'une population?"
+    ],
+    ["Qual è il principio di Le Chatelier e come si applica agli equilibri chimici?"],
+    [
+        "¿Qué es una supernova y cuál es su importancia en la formación de elementos pesados en el universo?"
+    ],
+    [
+        "Qual é a definição formal de uma integral de linha e como é utilizada em física?"
+    ],
+    [
+        "Was versteht man unter dem Moho-Diskontinuität und welche Bedeutung hat sie für das Verständnis der Erdkruste?"
+    ],
+    [
+        "Hiện tượng nhà kính là gì và nó ảnh hưởng như thế nào đến biến đổi khí hậu toàn cầu?"
+    ],
+    [
+        "알고리즘의 시간 복잡도가 중요한 이유는 무엇이며, 시간 복잡도를 어떻게 분석하나요?"
+    ],
+    ["什么是CRISPR-Cas9基因编辑技术，它在现代生物学研究中的作用是什么？"],
+    [
+        "Create a Python function that takes a list of integers and returns the list sorted in ascending order without using the built-in sort or sorted functions."
+    ],
+    [
+        "Écrivez une fonction en C++ qui trouve le plus long sous-tableau contigu avec une somme égale à zéro."
+    ],
+    [
+        "Scrivi una funzione in Java che calcola il fattoriale di un numero utilizzando la ricorsione."
+    ],
+    [
+        "Desarrolla una función en JavaScript que determine si una cadena de texto es un palíndromo, ignorando espacios y signos de puntuación."
+    ],
+    ["Implemente uma função em C# que verifique se uma matriz quadrada é simétrica."],
+    [
+        "Schreiben Sie eine Funktion in Swift, die eine gegebene Zeichenfolge in umgekehrter Reihenfolge zurückgibt, ohne integrierte Funktionen zu verwenden."
+    ],
+    [
+        "Viết một hàm trong PHP để tìm tất cả các số nguyên tố trong một khoảng cho trước."
+    ],
+    [
+        "파이썬을 사용하여 주어진 이진 트리가 이진 탐색 트리인지 확인하는 함수를 작성하십시오."
+    ],
+    [
+        "用 Go 语言编写一个函数，计算给定字符串中每个字符出现的次数，并返回一个包含字符及其出现次数的映射。"
+    ],
+    [
+        "Can you help me design a detailed project plan for developing a machine learning model for predicting stock prices?"
+    ],
+    [
+        "Pouvez-vous m'aider à organiser un emploi du temps hebdomadaire pour maximiser la productivité de mon équipe de développement logiciel?"
+    ],
+    [
+        "Puoi aiutarmi a creare un piano di sviluppo per un'applicazione mobile che gestisce le prenotazioni di ristoranti?"
+    ],
+    [
+        "¿Podrías ayudarme a elaborar un plan detallado para la implementación de un sistema de gestión de contenido (CMS) en una empresa mediana?"
+    ],
+    [
+        "Você pode me ajudar a planejar uma estratégia de desenvolvimento para um sistema de comércio eletrônico escalável?"
+    ],
+    [
+        "Können Sie mir helfen, einen detaillierten Zeitplan für die Implementierung eines neuen ERP-Systems in unserem Unternehmen zu erstellen?"
+    ],
+    [
+        "Bạn có thể giúp tôi xây dựng một kế hoạch phát triển chi tiết cho dự án xây dựng hệ thống quản lý chuỗi cung ứng không?"
+    ],
+    [
+        "신경망 기반 이미지 인식 모델 개발을 위한 세부 프로젝트 계획을 세우는 데 도움을 줄 수 있나요?"
+    ],
+    ["你能帮我制定一个详细的开发计划，用于创建一个基于区块链的分布式账本系统吗？"],
+    [
+        "Prove that the sum of the squares of any two sides of a right triangle is equal to the square of the hypotenuse."
+    ],
+    [
+        "Calculez la force gravitationnelle entre deux masses de 10 kg chacune séparées par une distance de 1 mètre."
+    ],
+    [
+        "Determina la formula molecolare di un composto che contiene il 40% di carbonio, il 6.67% di idrogeno e il 53.33% di ossigeno in massa."
+    ],
+    [
+        "Explica la teoría del ciclo económico de Schumpeter y cómo se aplica a la economía moderna."
+    ],
+    [
+        "Calcule a energia potencial gravitacional de um objeto de 5 kg a uma altura de 10 metros acima do solo (g = 9,8 m/s²)."
+    ],
+    [
+        "Beweisen Sie, dass jede Primzahl der Form 4k+1 als Summe zweier Quadrate geschrieben werden kann."
+    ],
+    [
+        "Tính nồng độ mol của dung dịch H₂SO₄ khi hoà tan 98 gam H₂SO₄ vào nước để được 1 lít dung dịch."
+    ],
+    ["케인스 경제학의 핵심 개념과 그것이 현대 경제 정책에 미치는 영향을 설명하십시오."],
+    ["计算一个质量为2 kg的物体在3米高处的重力势能（g = 9.8 m/s²）。"],
+    [
+        'Identify the author of a novel that features a dystopian society where "Big Brother" watches over its citizens and the protagonist works for the Ministry of Truth.'
+    ],
+    [
+        "Quel est le seul mammifère capable de voler activement, souvent associé à la nuit et capable d'écholocalisation?"
+    ],
+    [
+        "Qual è l'opera letteraria italiana che narra il viaggio immaginario di un poeta attraverso Inferno, Purgatorio e Paradiso, guidato da Virgilio e Beatrice?"
+    ],
+    [
+        "¿Qué insecto es conocido por su organización social compleja, su capacidad para producir miel y su comunicación mediante la danza?"
+    ],
+    [
+        "Qual é o fenômeno atmosférico que ocorre quando uma massa de ar quente se encontra com uma massa de ar frio, resultando em uma violenta tempestade giratória?"
+    ],
+    [
+        "Welches literarische Werk beschreibt die Geschichte eines jungen Mädchens, das durch einen Kaninchenbau in eine fantastische Welt voller skurriler Charaktere fällt?"
+    ],
+    [
+        "Động vật nào có thể tái sinh toàn bộ cơ thể từ một mảnh nhỏ của chính nó, thường sống dưới nước và có thể có nhiều xúc tu?"
+    ],
+    [
+        "어떤 자연 현상은 태양빛이 대기 중의 물방울에 반사되고 굴절되어 발생하며, 하늘에 나타나는 여러 색깔의 아치 형태를 띠나요?"
+    ],
+    ["这部文学作品讲述了一位绅士和他的侍从的冒险故事，他们在"],
+    [
+        "Can you derive the Euler-Lagrange equation from the principle of stationary action in classical mechanics?"
+    ],
+    [
+        "Expliquez la notion de « différence ontologique » chez Martin Heidegger et son importance pour la phénoménologie."
+    ],
+    [
+        "Qual è il significato simbolico del colore blu nei dipinti di Giotto di Bondone durante il Rinascimento?"
+    ],
+    [
+        "¿Cómo afecta el cambio de código a la estructura gramatical en comunidades bilingües de habla español-inglés?"
+    ],
+    [
+        "Qual é o impacto da política monetária não convencional no controle da inflação durante uma crise econômica?"
+    ],
+    [
+        "Erklären Sie den Unterschied zwischen deterministischen und nicht-deterministischen endlichen Automaten und ihre Anwendungsbereiche."
+    ],
+    [
+        "Giải thích cơ chế của quá trình phiên mã ngược (reverse transcription) và tầm quan trọng của nó trong nghiên cứu HIV/AIDS."
+    ],
+    ["조선시대 성리학이 한국 사회와 문화에 미친 영향을 설명하세요."],
+    ["如何解释量子纠缠现象，以及它在量子计算中的潜在应用？"],
+    [
+        "How can you design a daily schedule that maximizes productivity for a remote worker who has multiple meetings and project deadlines?"
+    ],
+    [
+        "Quels sont les meilleures stratégies pour gérer les conflits au sein d'une équipe multiculturelle travaillant sur un projet commun?"
+    ],
+    [
+        "Quali sono i migliori consigli per mantenere un equilibrio tra vita professionale e vita privata in un ambiente lavorativo stressante?"
+    ],
+    [
+        "¿Cómo se puede elaborar un plan financiero personal efectivo que incluya ahorro para la jubilación, inversión y manejo de deudas?"
+    ],
+    [
+        "Quais são as melhores práticas para implementar metodologias ágeis em uma equipe de desenvolvimento de software?"
+    ],
+    [
+        "Welche Strategien können verwendet werden, um ein starkes berufliches Netzwerk aufzubauen und zu pflegen, insbesondere in der Tech-Branche?"
+    ],
+    [
+        "Những bước nào cần thiết để xây dựng một lộ trình phát triển sự nghiệp bền vững trong lĩnh vực công nghệ thông tin?"
+    ],
+    ["프로젝트의 범위 변동을 효과적으로 관리하기 위한 최고의 방법은 무엇인가요?"],
+    ["在快速变化的职场环境中，如何有效地实现工作与生活的平衡？"],
+    [
+        "Write an argumentative essay discussing the pros and cons of artificial intelligence in the workplace, including potential ethical concerns."
+    ],
+    [
+        "Analysez les impacts sociaux et économiques de la digitalisation sur les petites entreprises en France."
+    ],
+    [
+        "Scrivi un'email formale al direttore di una rivista per proporre un articolo sulla sostenibilità ambientale nelle città italiane."
+    ],
+    [
+        "Elabora un informe detallado sobre los efectos del cambio climático en la biodiversidad de la región amazónica."
+    ],
+    [
+        "Analise criticamente os principais pontos abordados no relatório anual do Banco Mundial sobre a pobreza global."
+    ],
+    [
+        "Erstellen Sie eine technische Dokumentation für die Implementierung eines neuen Software-Features in einer bestehenden Anwendung."
+    ],
+    [
+        "Viết một bài luận phân tích về tác động của cuộc cách mạng công nghiệp 4.0 đối với thị trường lao động Việt Nam."
+    ],
+    [
+        "인공지능의 윤리적 문제에 대한 연구 논문을 작성하고, 다양한 사례를 통해 그 영향을 분석하세요."
+    ],
+    ["分析鲁迅的小说《阿Q正传》中反映的中国社会问题和作者的批判态度。"],
+]
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+
+if torch.cuda.is_available():
+    model_id = "lamhieu/ghost-8b-beta-disl-0x5-8k"
+    model_tk = os.getenv("HF_TOKEN", None)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        trust_remote_code=True,
+        token=model_tk,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        token=model_tk,
+    )
+    SelfExtend.apply(
+        model,
+        group_size=16,
+        window_size=512,
+        enable_flash_attention=True,
+        flash_attention_impl="flash_attn",
+    )
+    model.generation_config.max_length = 123392
+
+
+@spaces.GPU(duration=120)
+def generate(
+    message: str,
+    chat_history: list[tuple[str, str]],
+    system_prompt: str,
+    max_new_tokens: int = 1536,
+    temperature: float = 0.4,
+    top_p: float = 0.95,
+    top_k: int = 50,
+    repetition_penalty: float = 1.0,
+) -> Iterator[str]:
+    conversation = []
+    if system_prompt:
+        conversation.append({"role": "system", "content": system_prompt})
+    for user, assistant in chat_history:
+        conversation.extend(
+            [
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant},
+            ]
+        )
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(
+        conversation, add_generation_prompt=True, return_tensors="pt"
+    )
+    input_ids = input_ids.to(model.device)
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(
+            f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens."
+        )
+
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
+
+
+chatbot = gr.Chatbot(height=500, placeholder=PLACEHOLDER, label="Ghost 8B Beta")
+
+chat_interface = gr.ChatInterface(
+    fn=generate,
+    chatbot=chatbot,
+    fill_height=True,
+    additional_inputs=[
+        gr.Textbox(label="System prompt", lines=6),
+        gr.Slider(
+            label="Max new tokens",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=2.0,
+            step=0.1,
+            value=0.4,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.95,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=100,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.0,
+        ),
+    ],
+    stop_btn="Stop",
+    cache_examples=False,
+    examples=EXAMPLES,
+    examples_per_page=9,
+)
+
+with gr.Blocks(fill_height=True, css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    chat_interface.render()
+    gr.Markdown(LICENSE)
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(share=True)
+    # demo.launch(share=True)

requirements.txt

@@ -0,0 +1,8 @@
+accelerate==0.30.1
+bitsandbytes==0.43.1
+gradio==4.37.2
+scipy==1.13.0
+sentencepiece==0.2.0
+spaces==0.28.3
+torch==2.0.0
+transformers==4.41.0

self_extend_patch/Llama.py

@@ -0,0 +1,482 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import math
+from typing import Optional, Tuple
+import torch.nn.functional as F
+from transformers.cache_utils import Cache
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+from .selfextend_flash_attn import self_extend_flash_forward 
+
+
+
+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)
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin) if not q is None else None
+    k_embed = (k * cos) + (rotate_half(k) * sin) if not k is None else None
+    return q_embed, k_embed
+
+
+
+
+def self_extend_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Cache] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+    cache_position: Optional[torch.LongTensor] = None,
+    group_size_1: Optional[float] = 8,
+    group_size_2: Optional[float] = 1024,
+    scale_base: Optional[int] = -1,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    if "padding_mask" in kwargs:
+        warnings.warn(
+            "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+        )
+
+    bsz, q_len, _ = hidden_states.size()
+
+
+    if self.config.pretraining_tp > 1:
+        key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
+        query_slices = self.q_proj.weight.split(
+            (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
+        )
+        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+        query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
+        query_states = torch.cat(query_states, dim=-1)
+
+        key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
+        key_states = torch.cat(key_states, dim=-1)
+
+        value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
+        value_states = torch.cat(value_states, dim=-1)
+
+    else:
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    if scale_base > 0:
+        scaled_query = query_states * ((position_ids + 1)[:, None, :, None].log() / np.log(scale_base)).clip(1).to(query_states.dtype) # log scale 
+        #scaled_query = query_states * (((0.1*(((position_ids+1)[:, None, :, None]/scale_base).log())+1)**2).clip(1)).to(query_states.dtype) # Yarn scale 
+    else:
+        scaled_query = query_states
+    
+    past_key_value = getattr(self, "past_key_value", past_key_value)
+    if past_key_value is not None:
+        # sin and cos are specific to RoPE models; position_ids needed for the static cache
+        cache_kwargs = {"cache_position": cache_position}
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    kv_seq_len = key_states.shape[-2]
+
+    query_position = position_ids
+    key_position = position_ids if q_len != 1 else torch.arange(kv_seq_len, dtype=position_ids.dtype).to(query_position.device).view(1, kv_seq_len) # only consider bsz=1 for now.
+
+
+
+    neighbor_q_cos, neighbor_q_sin = self.rotary_emb(value_states, query_position)#, seq_len=None)
+    neighbor_k_cos, neighbor_k_sin = self.rotary_emb(value_states, key_position)#, seq_len=None)
+
+
+    _re_group_size_2 = 0 if query_position.max() < group_size_2 else group_size_2 # in case that, the smallest q position, g2-g2//g1 exceed the max position
+    group_query_position = query_position // group_size_1 + _re_group_size_2 - _re_group_size_2 // group_size_1
+    group_key_position = key_position // group_size_1
+
+    group_q_cos, group_q_sin = self.rotary_emb(value_states, group_query_position)#, seq_len=None)
+    group_k_cos, group_k_sin = self.rotary_emb(value_states, group_key_position)#, seq_len=None)
+
+
+
+    neighbor_query_states, _ = apply_rotary_pos_emb(scaled_query, None, neighbor_q_cos, neighbor_q_sin, None)
+    _, neighbor_key_states = apply_rotary_pos_emb(None, key_states, neighbor_k_cos, neighbor_k_sin, None)
+    group_query_states, _ = apply_rotary_pos_emb(scaled_query, None, group_q_cos, group_q_sin, None)
+    _, group_key_states = apply_rotary_pos_emb(None, key_states, group_k_cos, group_k_sin, None)
+
+
+
+    neighbor_key_states = repeat_kv(neighbor_key_states, self.num_key_value_groups)
+    group_key_states = repeat_kv(group_key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+
+
+    neighbor_attn_weights = torch.matmul(neighbor_query_states, neighbor_key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+    group_attn_weights = torch.matmul(group_query_states, group_key_states.transpose(2, 3)) / math.sqrt(self.head_dim) 
+
+
+    if attention_mask is not None:  # no matter the length, we just slice it
+        if cache_position is not None:
+            causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]]
+        else:
+            causal_mask = attention_mask
+        group_attn_weights = group_attn_weights + causal_mask
+        neighbor_attn_weights = neighbor_attn_weights + causal_mask
+
+    if q_len == 1:
+        neighbor_attention_mask = torch.zeros((q_len, kv_seq_len), device=neighbor_attn_weights.device)
+        neighbor_attention_mask[:, -group_size_2:] = 1
+    elif q_len == kv_seq_len:
+        neighbor_attention_mask = torch.ones((q_len, kv_seq_len), device=neighbor_attn_weights.device)
+        neighbor_attention_mask = torch.tril(neighbor_attention_mask)
+        if q_len-group_size_2 > 0:
+            group_attention_mask =  torch.tril(torch.ones((q_len-group_size_2, kv_seq_len-group_size_2), device=group_attn_weights.device))
+            neighbor_attention_mask[group_size_2:, :-group_size_2] -= group_attention_mask
+    else:
+        raise ValueError("q_len should be 1 or seq_len.")
+    
+    neighbor_attention_mask = neighbor_attention_mask.bool()
+    attn_weights = torch.where(neighbor_attention_mask, neighbor_attn_weights, group_attn_weights)
+    
+    # upcast attention to fp32
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+
+    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+        raise ValueError(
+            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+            f" {attn_output.size()}"
+        )
+    
+
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+    if self.config.pretraining_tp > 1:
+        attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
+        o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
+        attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
+    else:
+        attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+
+    return attn_output, attn_weights, past_key_value
+
+def flash_self_extend_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Cache] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+    group_size_1: Optional[float] = 8,
+    group_size_2: Optional[float] = 1024,
+    scale_base: Optional[int] = -1,
+    cache_position: Optional[torch.LongTensor] = None,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """
+        Require updating tansformers to >= 4.38.2, flash_attn >= 2.5.6
+        a. Only support causal mask.
+        b. Don't support atttention_mask.
+        c. Never test it with batch size > 1.
+        d. Only support q_len = 1 or q_len = seq_len.
+    """
+    if "padding_mask" in kwargs:
+        warnings.warn(
+            "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+        )
+        attention_mask = kwargs.pop("padding_mask")
+
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    if scale_base > 0:
+        scaled_query = query_states * ((position_ids + 1)[:, None, :, None].log() / np.log(scale_base)).clip(1).to(query_states.dtype) # log scale 
+        #scaled_query = query_states * (((0.1*(((position_ids+1)[:, None, :, None]/scale_base).log())+1)**2).clip(1)).to(query_states.dtype) # Yarn scale 
+    else:
+        scaled_query = query_states
+    
+    past_key_value = getattr(self, "past_key_value", past_key_value)
+    if past_key_value is not None:
+        # sin and cos are specific to RoPE models; position_ids needed for the static cache
+        cache_kwargs = {"cache_position": cache_position}
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    kv_seq_len = key_states.shape[-2]
+
+    query_position = position_ids
+    # only consider bsz=1 for now. 
+    key_position = position_ids if q_len != 1 else torch.arange(kv_seq_len, dtype=position_ids.dtype).to(query_position.device).view(1, kv_seq_len) 
+    attn_dropout = self.config.attention_dropout if self.training else 0.0
+    if q_len == 1:
+        # We implement the case q_len == 1 separately, by manipulating positions.
+        # for our flash implementation doesnot work for  decoding stage at the releasing time.
+
+        neighbor_key_position = position_ids[:, -1] - key_position
+        _re_group_size_2 = 0 if position_ids.max() < group_size_2 else group_size_2
+        group_key_position = position_ids[:, -1]//group_size_1 - key_position//group_size_1 + (_re_group_size_2 - _re_group_size_2//group_size_1)
+        decode_key_position = torch.cat([group_key_position[:, :-group_size_2], neighbor_key_position[:,-group_size_2:]], dim=1)
+        
+        decode_k_cos, decode_k_sin = self.rotary_emb(value_states, decode_key_position)#, seq_len=None)
+        #import pdb; pdb.set_trace()
+        #neighbor_query_states, _ = apply_rotary_pos_emb(scaled_query, None, cos, sin, query_position_ids) 
+        decode_query_states = scaled_query.transpose(1,2).contiguous() # position 0: cos 0 = 1, sin 0 = 0
+        _, decode_key_states = apply_rotary_pos_emb(None, key_states, decode_k_cos, -decode_k_sin, decode_key_position) 
+
+        decode_key_states = repeat_kv(decode_key_states, self.num_key_value_groups).transpose(1, 2).contiguous()
+        decode_value_states = repeat_kv(value_states, self.num_key_value_groups).transpose(1, 2).contiguous()
+        
+        attn_output = flash_attn_func(decode_query_states,
+                                      decode_key_states,
+                                      decode_value_states,
+                                      attn_dropout, 
+                                      softmax_scale=None, 
+                                      causal=True)
+    elif q_len == kv_seq_len:
+        # set correct position_ids & apply RoPE.
+        neighbor_q_cos, neighbor_q_sin = self.rotary_emb(value_states, query_position)#, seq_len=None)
+        neighbor_k_cos, neighbor_k_sin = self.rotary_emb(value_states, key_position)#, seq_len=None)
+
+        _re_group_size_2 = 0 if query_position.max() < group_size_2 else group_size_2 # in case that, the smallest q position, g2-g2//g1 exceed the max position
+        group_query_position = query_position // group_size_1 + _re_group_size_2 - _re_group_size_2 / group_size_1
+        group_key_position = key_position // group_size_1
+
+        group_q_cos, group_q_sin = self.rotary_emb(value_states, group_query_position)#, seq_len=None)
+        group_k_cos, group_k_sin = self.rotary_emb(value_states, group_key_position)#, seq_len=None)
+
+        neighbor_query_states, _ = apply_rotary_pos_emb(scaled_query, None, neighbor_q_cos, neighbor_q_sin, None)
+        _, neighbor_key_states = apply_rotary_pos_emb(None, key_states, neighbor_k_cos, neighbor_k_sin, None)
+        group_query_states, _ = apply_rotary_pos_emb(scaled_query, None, group_q_cos, group_q_sin, None)
+        _, group_key_states = apply_rotary_pos_emb(None, key_states, group_k_cos, group_k_sin, None)
+        
+
+        neighbor_query_states = neighbor_query_states.transpose(1, 2).contiguous()
+        neighbor_key_states = repeat_kv(neighbor_key_states, self.num_key_value_groups).transpose(1, 2).contiguous()
+        group_query_states = group_query_states.transpose(1, 2).contiguous()
+        group_key_states = repeat_kv(group_key_states, self.num_key_value_groups).transpose(1, 2).contiguous()
+        value_states = repeat_kv(value_states, self.num_key_value_groups).transpose(1, 2).contiguous()
+
+        attn_output = self_extend_flash_forward(self,
+                                                query_position,
+                                                group_size_2,
+                                                neighbor_query_states,
+                                                neighbor_key_states,
+                                                group_query_states,
+                                                group_key_states,
+                                                value_states,
+                                                attention_mask,
+                                                bsz,
+                                                q_len,
+                                                kv_seq_len,
+                                                attn_dropout,
+                                            )
+    else:
+        raise ValueError("q_len should be 1 or seq_len.")
+    
+    attn_output = attn_output.contiguous()
+    attn_output = attn_output.view(bsz, q_len, -1).contiguous()
+    attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+    return attn_output, attn_weights, past_key_value
+
+
+
+def lm_infinite_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Cache] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+    cache_position: Optional[torch.LongTensor] = None,
+    group_size_1: Optional[float] = 8,
+    group_size_2: Optional[float] = 1024,
+    initial_num: Optional[int] = 1,
+    scale_base: Optional[int] = -1,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    if "padding_mask" in kwargs:
+        warnings.warn(
+            "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+        )
+
+    bsz, q_len, _ = hidden_states.size()
+
+
+    if self.config.pretraining_tp > 1:
+        key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
+        query_slices = self.q_proj.weight.split(
+            (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
+        )
+        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+        query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
+        query_states = torch.cat(query_states, dim=-1)
+
+        key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
+        key_states = torch.cat(key_states, dim=-1)
+
+        value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
+        value_states = torch.cat(value_states, dim=-1)
+
+    else:
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    if scale_base > 0:
+        scaled_query = query_states * ((position_ids + 1)[:, None, :, None].log() / np.log(scale_base)).clip(1).to(query_states.dtype) # log scale 
+        #scaled_query = query_states * (((0.1*(((position_ids+1)[:, None, :, None]/scale_base).log())+1)**2).clip(1)).to(query_states.dtype) # Yarn scale 
+    else:
+        scaled_query = query_states
+    
+    past_key_value = getattr(self, "past_key_value", past_key_value)
+    if past_key_value is not None:
+        # sin and cos are specific to RoPE models; position_ids needed for the static cache
+        cache_kwargs = {"cache_position": cache_position}
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    kv_seq_len = key_states.shape[-2]
+
+    query_position = position_ids
+    key_position = position_ids if q_len != 1 else torch.arange(kv_seq_len, dtype=position_ids.dtype).to(query_position.device).view(1, kv_seq_len) # only consider bsz=1 for now.
+
+
+
+    neighbor_q_cos, neighbor_q_sin = self.rotary_emb(value_states, query_position)#, seq_len=None)
+    neighbor_k_cos, neighbor_k_sin = self.rotary_emb(value_states, key_position)#, seq_len=None)
+
+
+    _re_group_size_2 = 0 if query_position.max() < group_size_2 else group_size_2 # in case that, the smallest q position, g2-g2//g1 exceed the max position
+    group_query_position = query_position // group_size_1 + _re_group_size_2 - _re_group_size_2 / group_size_1
+    group_key_position = key_position // group_size_1
+
+    group_q_cos, group_q_sin = self.rotary_emb(value_states, group_query_position)#, seq_len=None)
+    group_k_cos, group_k_sin = self.rotary_emb(value_states, group_key_position)#, seq_len=None)
+
+
+
+    neighbor_query_states, _ = apply_rotary_pos_emb(scaled_query, None, neighbor_q_cos, neighbor_q_sin, None)
+    _, neighbor_key_states = apply_rotary_pos_emb(None, key_states, neighbor_k_cos, neighbor_k_sin, None)
+    group_query_states, _ = apply_rotary_pos_emb(scaled_query, None, group_q_cos, group_q_sin, None)
+    _, group_key_states = apply_rotary_pos_emb(None, key_states, group_k_cos, group_k_sin, None)
+
+
+
+    neighbor_key_states = repeat_kv(neighbor_key_states, self.num_key_value_groups)
+    group_key_states = repeat_kv(group_key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+
+
+    neighbor_attn_weights = torch.matmul(neighbor_query_states, neighbor_key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+    group_attn_weights = torch.matmul(group_query_states, group_key_states.transpose(2, 3)) / math.sqrt(self.head_dim) 
+
+
+    if attention_mask is not None:  # no matter the length, we just slice it
+        if cache_position is not None:
+            causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]]
+        else:
+            causal_mask = attention_mask
+        group_attn_weights = group_attn_weights + causal_mask
+        neighbor_attn_weights = neighbor_attn_weights + causal_mask
+
+    if q_len == 1:
+        neighbor_attention_mask = torch.zeros((q_len, kv_seq_len), device=neighbor_attn_weights.device)
+        neighbor_attention_mask[:, -group_size_2:] = 1
+    elif q_len == kv_seq_len:
+        neighbor_attention_mask = torch.ones((q_len, kv_seq_len), device=neighbor_attn_weights.device)
+        neighbor_attention_mask = torch.tril(neighbor_attention_mask)
+        if q_len-group_size_2 > 0:
+            group_attention_mask =  torch.tril(torch.ones((q_len-group_size_2, kv_seq_len-group_size_2), device=group_attn_weights.device))
+            neighbor_attention_mask[group_size_2:, :-group_size_2] -= group_attention_mask
+    else:
+        raise ValueError("q_len should be 1 or seq_len.")
+    
+    neighbor_attention_mask = neighbor_attention_mask.bool()
+    attn_weights = torch.where(neighbor_attention_mask, neighbor_attn_weights, group_attn_weights)
+    
+    # upcast attention to fp32
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+
+    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+        raise ValueError(
+            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+            f" {attn_output.size()}"
+        )
+    
+
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+    if self.config.pretraining_tp > 1:
+        attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
+        o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
+        attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
+    else:
+        attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+
+    return attn_output, attn_weights, past_key_value

self_extend_patch/__init__.py

@@ -0,0 +1 @@
+from . import Llama

self_extend_patch/selfextend_flash_attn.py

@@ -0,0 +1,199 @@
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+import torch
+
+# must replace orginal flash forward method with the following one first, to enbale the window feature.
+def flash_attention2_forward_with_window_size(
+    self,
+    query_states,
+    key_states,
+    value_states,
+    attention_mask,
+    query_length,
+    dropout=0.0,
+    softmax_scale=None,
+    window_size=[-1, -1],
+    return_attn_probs=False,
+):
+    """
+    Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+    first unpad the input, then computes the attention scores and pad the final attention scores.
+
+    Args:
+        query_states (`torch.Tensor`):
+            Input query states to be passed to Flash Attention API
+        key_states (`torch.Tensor`):
+            Input key states to be passed to Flash Attention API
+        value_states (`torch.Tensor`):
+            Input value states to be passed to Flash Attention API
+        attention_mask (`torch.Tensor`):
+            The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+            position of padding tokens and 1 for the position of non-padding tokens.
+        dropout (`int`, *optional*):
+            Attention dropout
+        softmax_scale (`float`, *optional*):
+            The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        window_size ([Int, Int])
+            The left & right window size for Flash Attention. Default to [-1, -1] which means no window size is used.
+        return_attn_probs (`bool`, *optional*):
+            Whether to return the attention softmax logssumexp and probabilities. Default to False.
+    """
+    if not self._flash_attn_uses_top_left_mask:
+        causal = self.is_causal
+    else:
+        # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+        causal = self.is_causal and query_length != 1
+
+    # Contains at least one padding token in the sequence
+    if attention_mask is not None:
+        batch_size = query_states.shape[0]
+        query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+            query_states, key_states, value_states, attention_mask, query_length
+        )
+
+        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+        max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+        attn_output_unpad, softmax_lse, S_dmask = flash_attn_varlen_func(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q,
+            max_seqlen_k=max_seqlen_in_batch_k,
+            dropout_p=dropout,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            return_attn_probs=True,
+        )
+        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+    else:
+        attn_output, softmax_lse, S_dmask = flash_attn_func(
+            query_states,
+            key_states,
+            value_states,
+            dropout,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            return_attn_probs=True,
+        )
+
+    if return_attn_probs:
+        return attn_output, softmax_lse, S_dmask
+    else:
+        return attn_output
+
+def self_extend_flash_forward(
+        model_self,
+        query_position,
+        group_size_2,
+        neighbor_query_states,
+        neighbor_key_states,
+        group_query_states,
+        group_key_states,
+        value_states,
+        attention_mask,
+        bsz,
+        q_len,
+        kv_seq_len,
+        attn_dropout,
+    ):
+    
+    if query_position.max() >= group_size_2:
+        neighbor_attn_output, neighbor_softmax_lse_right_padded, neighbor_prob = model_self._flash_attention_forward(
+            neighbor_query_states,
+            neighbor_key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=attn_dropout,
+            window_size=[group_size_2 - 1, 0],
+            # right dim here does not matter and can be -1, or > 0 due to causal mask
+            return_attn_probs=True,
+        )
+
+        group_attention_len = (
+            kv_seq_len - group_size_2
+        )  # here we should use kv_seq_len rather than max_kv_len since we have paddings in qkv and attention_mask
+
+        group_attention_mask = attention_mask[:, :group_attention_len] if not attention_mask is None else None
+        group_attn_output, group_softmax_lse_right_padded, group_prob = model_self._flash_attention_forward(
+            group_query_states[:, -group_attention_len:, :, :],
+            group_key_states[:, :group_attention_len, :, :],
+            value_states[:, :group_attention_len, :, :],
+            group_attention_mask,
+            group_query_states[:, -group_attention_len:, :, :].shape[1],
+            dropout=attn_dropout,
+            window_size=[-1, -1],
+            return_attn_probs=True,
+        )  # note that kv and q's indexing are different! also query size could be different from kv length and very small during generation compared to prefilling
+
+
+        # normalize lse first
+        neighbor_seq_length = torch.Tensor([kv_seq_len,]).long().expand(bsz, 1) if attention_mask is None else torch.sum(attention_mask, axis=1, keepdim=True)  # [batch_size, 1]
+        group_seq_length = torch.Tensor([group_attention_len,]).long().expand(bsz, 1) if attention_mask is None else torch.sum(attention_mask[:, :group_attention_len], axis=1, keepdim=True)  # [batch_size, 1]
+
+        # convert align left to align right and convert exp(0) to 0
+        neighbor_softmax_lse = torch.zeros_like(neighbor_softmax_lse_right_padded)
+        group_softmax_lse = torch.zeros_like(group_softmax_lse_right_padded)
+        for idx in range(bsz):
+            if neighbor_seq_length[idx] > 0:
+                neighbor_softmax_lse[idx, :, -neighbor_seq_length[idx] :] = neighbor_softmax_lse_right_padded[
+                    idx, :, : neighbor_seq_length[idx]
+                ]
+            if group_seq_length[idx] > 0:
+                group_softmax_lse[idx, :, -group_seq_length[idx] :] = group_softmax_lse_right_padded[
+                    idx, :, : group_seq_length[idx]
+                ]
+
+        # attn_output size is [batch_size, max_seq_len (not the true one), query_length, dim]
+        true_neighbor_seq_max_length = neighbor_softmax_lse.shape[
+            -1
+        ]  # it could be smaller than query_length due to the attention_mask
+        true_group_seq_max_length = group_softmax_lse.shape[
+            -1
+        ]  # it could be smaller than group_query_layer[:, -group_attention_len:, :, :].shape[1] due to the attention_mask[:, :group_attention_len]
+
+        neighbor_softmax_lse = neighbor_softmax_lse.transpose(1, 2).unsqueeze(
+            -1
+        )  # [batch_size, true_neighbor_seq_max_length, self.num_heads, 1]
+        group_softmax_lse = group_softmax_lse.transpose(1, 2).unsqueeze(
+            -1
+        )  # [batch_size, true_group_seq_max_length, self.num_heads, 1]
+
+        lse_gap = group_softmax_lse - neighbor_softmax_lse[:, -true_group_seq_max_length:, :, :]
+        #if  torch.isinf(neighbor_softmax_lse).any() or torch.isnan(neighbor_softmax_lse).any():
+        #    import pdb; pdb.set_trace()
+        
+        neighbor_softmax_lse[:, -true_group_seq_max_length:, :, :] = 1 / (1 + torch.exp(lse_gap))
+        neighbor_softmax_lse[:, :-true_group_seq_max_length, :, :] = 1.
+        group_softmax_lse = 1 / (1 + torch.exp(-lse_gap))
+
+
+
+        neighbor_attn_output[:, -true_neighbor_seq_max_length:, ...] = (
+            neighbor_attn_output[:, -true_neighbor_seq_max_length:, ...] * neighbor_softmax_lse
+        )
+        group_attn_output[:, -true_group_seq_max_length:, ...] = (
+            group_attn_output[:, -true_group_seq_max_length:, ...] * group_softmax_lse
+        )
+        attn_output = torch.empty_like(neighbor_attn_output).copy_(
+            neighbor_attn_output
+        )  # might be slightly faster than clone
+        #attn_output[:, group_size_2:, ...] += group_attn_output
+        attn_output[:, group_size_2-kv_seq_len:, ...] += group_attn_output
+        attn_output = torch.nan_to_num(attn_output, nan=0)  
+    
+    else:
+        attn_output = model_self._flash_attention_forward(
+            neighbor_query_states,
+            neighbor_key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=attn_dropout,
+            window_size=[-1, -1],
+        )
+
+    return attn_output

self_extend_patch/selfextend_flash_attn_triton.py

@@ -0,0 +1,278 @@
+import math
+import torch
+
+import triton
+import triton.language as tl
+
+
+def self_extend_flash_forward_triton(
+        model_self,
+        query_position,
+        group_size_2,
+        neighbor_query_states,
+        neighbor_key_states,
+        group_query_states,
+        group_key_states,
+        value_states,
+        attention_mask,
+        bsz,
+        q_len,
+        kv_seq_len,
+        attn_dropout,
+    ):
+
+    o = _self_extend_flash_forward_triton(q=neighbor_query_states,
+                        k=neighbor_key_states,
+                        q1=group_query_states,
+                        k1=group_key_states,
+                        v=value_states,
+                        causal=(q_len == kv_seq_len),
+                        sm_scale=1. / math.sqrt(neighbor_query_states.shape[-1]),
+                        window=group_size_2)
+    o = o.transpose(1, 2).contiguous()
+    # print("o", o.shape)
+    return o
+
+
+
+
+
+
+def _self_extend_flash_forward_triton(q, k, q1, k1, v, causal, sm_scale, window):
+        # shape constraints
+        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+        assert Lq == Lk and Lk == Lv
+        assert Lk in {16, 32, 64, 128}
+        
+        device = torch.cuda.device_of(q)
+        with torch.cuda.device(device):
+            o = torch.empty_like(q)
+            BLOCK_M = 128
+            BLOCK_N = 32
+            grid = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1])
+            L = torch.empty((q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+            _fwd_kernel[grid](
+                q, 
+                k,
+                q1,
+                k1,
+                v, 
+                sm_scale,
+                L,
+                o,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+                q.shape[0], 
+                q.shape[1], 
+                q.shape[2],
+                k.shape[2],
+                BLOCK_M=BLOCK_M, 
+                BLOCK_N=BLOCK_N, 
+                BLOCK_DMODEL=Lk,
+                IS_CAUSAL=causal,
+                WINDOW=window,
+                num_warps=8,
+                num_stages=2)
+
+        return o
+
+
+
+
+@triton.heuristics(
+    {
+        "EVEN_M": lambda args: args["Q_CTX"] % args["BLOCK_M"] == 0,
+        "EVEN_N": lambda args: args["KV_CTX"] % args["BLOCK_N"] == 0,
+    }
+)
+@triton.jit
+def _fwd_kernel(
+    Q, 
+    K,
+    Q1,
+    K1,
+    V, 
+    sm_scale,
+    L,
+    Out,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_oz, stride_oh, stride_om, stride_on,
+    Z, 
+    H, 
+    Q_CTX,
+    KV_CTX,
+    BLOCK_M: tl.constexpr, 
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    WINDOW: tl.constexpr,
+    EVEN_M: tl.constexpr,
+    EVEN_N: tl.constexpr
+
+):
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    # qvk_offset = off_hz * stride_qh
+    q_offset = off_hz * stride_qh
+    vk_offset = off_hz * stride_kh
+    # vk_offset = q_offset
+
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    K_block_ptr = tl.make_block_ptr(
+        base=K + vk_offset,
+        shape=(KV_CTX, BLOCK_DMODEL),
+        strides=(stride_kn, stride_kk),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    Q1_block_ptr = tl.make_block_ptr(
+        base=Q1 + q_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    K1_block_ptr = tl.make_block_ptr(
+        base=K1 + vk_offset,
+        shape=(KV_CTX, BLOCK_DMODEL),
+        strides=(stride_kn, stride_kk),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=V + vk_offset,
+        shape=(KV_CTX, BLOCK_DMODEL),
+        strides=(stride_vn, stride_vk),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    qk_scale = sm_scale * 1.4426950408889634
+    
+    # load q: it will stay in SRAM throughout
+    if EVEN_M:
+        q = tl.load(Q_block_ptr)
+        q1 = tl.load(Q1_block_ptr)
+    else:
+        q = tl.load(Q_block_ptr, boundary_check=(1,0))
+        q1 = tl.load(Q1_block_ptr, boundary_check=(1,0))
+
+    q = (q * qk_scale).to(tl.bfloat16)
+    q1 = (q1 * qk_scale).to(tl.bfloat16)
+
+
+    # Dot I trick: it converts q1, q2 into mma layout and saves shared memory
+    # better way to generate a eye matrix. avoid casting from bool
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+    I = tl.where(offs_k[:, None] == offs_k,
+                 tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 1.0, dtype=tl.bfloat16),
+                 tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 0.0, dtype=tl.bfloat16))
+    q = tl.dot(q, I).to(tl.bfloat16)
+    q1 = tl.dot(q1, I).to(tl.bfloat16)
+
+
+    # loop over k, v and update accumulator
+    lo = 0
+    if IS_CAUSAL:
+        hi = tl.minimum(KV_CTX, (start_m + 1) * BLOCK_M)
+    else:
+        hi = KV_CTX
+
+    for start_n in range(lo, hi, BLOCK_N):
+        # -- load k, v --
+        if EVEN_N:
+            k = tl.load(K_block_ptr)
+            k1 = tl.load(K1_block_ptr)
+            v = tl.load(V_block_ptr)
+        else:
+            k = tl.load(K_block_ptr, boundary_check=(1,0))
+            k1 = tl.load(K1_block_ptr, boundary_check=(1,0))
+            v = tl.load(V_block_ptr, boundary_check=(1,0))
+        
+        # -- compute qk ---
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+
+        # Window masking
+        mask = ( KV_CTX - Q_CTX + offs_m[:, None]) >= (start_n + offs_n[None, :] + WINDOW)
+        qk += tl.where(mask, tl.dot(q1, tl.trans(k1)), tl.dot(q, tl.trans(k)))
+
+        # if not EVEN_N:
+        #     mask = (start_n + offs_n) < KV_CTX
+        #     qk = tl.where(mask, qk, float("-inf"))
+
+        if IS_CAUSAL:
+            mask = offs_m[:, None] >= (start_n + offs_n[None, :])
+            qk = tl.where(mask, qk, float("-inf"))
+        # qk += tl.dot(q, k)
+        
+        # -- compute scaling constant ---
+        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
+        alpha = tl.math.exp2(m_i - m_i_new)
+        p = tl.math.exp2(qk - m_i_new[:, None])
+        
+        # -- scale and update acc --
+        acc_scale = l_i * 0 + alpha  # workaround some compiler bug
+        acc *= acc_scale[:, None]
+        acc += tl.dot(p.to(tl.bfloat16), v)
+        
+        # -- update m_i and l_i --
+        l_i = l_i * alpha + tl.sum(p, 1)
+        m_i = m_i_new
+        
+        # update pointers
+        K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0))
+        K1_block_ptr = tl.advance(K1_block_ptr, (BLOCK_N, 0))
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+    
+
+    # write back l and m        
+    acc = acc * (1.0 / l_i[:, None])
+    l_ptrs = L + off_hz * Q_CTX + offs_m
+    
+    mask_m = offs_m < Q_CTX
+    l_i = m_i  + tl.math.log2(l_i)
+    if EVEN_M:
+        tl.store(l_ptrs, l_i)
+    else:
+        tl.store(l_ptrs, l_i, mask=mask_m)
+
+    # write back O
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + q_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0)
+    )
+    if EVEN_M:
+        tl.store(O_block_ptr, acc.to(tl.bfloat16))
+    else:
+        tl.store(O_block_ptr, acc.to(tl.bfloat16), boundary_check=(1,0))

self_extend_patch/triton_selfextend_flash_attn.py

@@ -0,0 +1,250 @@
+import torch
+import math
+import triton
+import triton.language as tl
+
+
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Uncommenting
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+#@triton.autotune(
+#    configs=[
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=4, num_warps=8),
+#        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64}, num_stages=3, num_warps=8),
+#        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 32}, num_stages=3, num_warps=8),
+#        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 32}, num_stages=3, num_warps=4),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=3, num_warps=4),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=4, num_warps=4),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=3, num_warps=4),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=4, num_warps=4),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=3, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=7, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=7, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=6, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=5, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_stages=4, num_warps=8),
+#        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_stages=6, num_warps=4),
+#    ],
+#    key=['N_CTX'],
+#)
+@triton.jit
+def _attn_fwd_prefill(Q1, K1, Q2, K2, V, sm_scale, M, Out,  #
+              stride_qz, stride_qh, stride_qm, stride_qk,  #
+              stride_kz, stride_kh, stride_kn, stride_kk,  #
+              stride_vz, stride_vh, stride_vk, stride_vn,  #
+              stride_oz, stride_oh, stride_om, stride_on,  #
+              Z, H,  #
+              Q_CTX: tl.constexpr,  #
+              N_CTX: tl.constexpr,  #
+              WINDOW: tl.constexpr,  #
+              BLOCK_M: tl.constexpr,  #
+              BLOCK_DMODEL: tl.constexpr,  #
+              BLOCK_N: tl.constexpr,  #
+              ):
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+
+    # block pointers
+    Q1_block_ptr = tl.make_block_ptr(
+        base=Q1 + qvk_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    Q2_block_ptr = tl.make_block_ptr(
+        base=Q2 + qvk_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=V + qvk_offset,
+        shape=(N_CTX, BLOCK_DMODEL),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    K1_block_ptr = tl.make_block_ptr(
+        base=K1 + qvk_offset,
+        shape=(BLOCK_DMODEL, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_DMODEL, BLOCK_N),
+        order=(0, 1),
+    )
+    K2_block_ptr = tl.make_block_ptr(
+        base=K2 + qvk_offset,
+        shape=(BLOCK_DMODEL, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_DMODEL, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + qvk_offset,
+        shape=(Q_CTX, BLOCK_DMODEL),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.442695040888963#1.44269504  # 1/log(2)
+    # load q: it will stay in SRAM throughout
+    #q = tl.load(Q_block_ptr)
+    if start_m * BLOCK_M + BLOCK_M > Q_CTX:
+        q1 = tl.load(Q1_block_ptr, boundary_check=(0,), padding_option='zero')
+        q2 = tl.load(Q2_block_ptr, boundary_check=(0,), padding_option='zero')
+    else:
+        q1 = tl.load(Q1_block_ptr)
+        q2 = tl.load(Q2_block_ptr)
+    #q1 = (q1 * qk_scale).to(tl.float16)
+    #q2 = (q2 * qk_scale).to(tl.float16)
+                
+    lo = 0
+    hi = (start_m + 1) * BLOCK_M 
+    # loop over k, v and update accumulator
+    for start_n in range(lo, hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) #?
+        #qk = qk.to(tl.float16)
+        # if use condition, qk has to be float32, then convert to float16...
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        if start_n + BLOCK_N - 1 > start_m * BLOCK_M - 1:
+            qk += tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), 0, -1.0e6)#float("-inf"))
+        
+        #qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
+        # -- compute qk ----
+        #k = tl.load(K_block_ptr)
+        # case 1: only need group attention: q2, k2
+        if BLOCK_N + start_n <= (start_m * BLOCK_M - WINDOW + 1):
+            if BLOCK_N + start_n >= N_CTX:
+                k2 = tl.load(K2_block_ptr, boundary_check=(1,), padding_option='zero')
+                v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero')
+            else:
+                k2 = tl.load(K2_block_ptr)
+                v = tl.load(V_block_ptr)
+            #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float16)
+            qk += tl.dot(q2, k2)#, out_dtype=tl.float16)
+        else:
+            #case 2: only need neighbor attention: q1, k1
+            if start_n >= (start_m+1) * BLOCK_M - WINDOW:
+                if BLOCK_N + start_n >= N_CTX:
+                    k1 = tl.load(K1_block_ptr, boundary_check=(1,), padding_option='zero')
+                    v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero')
+                else:
+                    k1 = tl.load(K1_block_ptr)
+                    v = tl.load(V_block_ptr)
+                #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+                #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float16)
+                qk += tl.dot(q1, k1)#, out_dtype=tl.float16)
+            else:
+                #case 3: need both q1, k1 and q2, k2
+                if BLOCK_N + start_n >= N_CTX:
+                    k1 = tl.load(K1_block_ptr, boundary_check=(1,), padding_option='zero')
+                    k2 = tl.load(K2_block_ptr, boundary_check=(1,), padding_option='zero')
+                    v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero')
+                else:
+                    k1 = tl.load(K1_block_ptr)
+                    k2 = tl.load(K2_block_ptr)
+                    v = tl.load(V_block_ptr)
+                #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+                #qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float16)
+                qk1 = tl.dot(q1, k1)#, out_dtype=tl.float16)
+                qk2 = tl.dot(q2, k2)#, out_dtype=tl.float16)
+                #merge_mask = tl.abs((offs_m[:, None] - (start_n + offs_n[None, :]))) >= WINDOW
+                #qk += tl.where(merge_mask, qk2, qk1)
+                qk += tl.where(tl.abs(offs_m[:, None] - (start_n + offs_n[None, :])) < WINDOW, qk1, qk2)
+
+        qk *= qk_scale
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1))
+        qk = qk - m_ij[:, None]
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        #v = tl.load(V_block_ptr)
+        #v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero')
+        acc += tl.dot(p.to(tl.float16), v)
+        # update m_i and l_i
+        m_i = m_ij
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K1_block_ptr = tl.advance(K1_block_ptr, (0, BLOCK_N))
+        K2_block_ptr = tl.advance(K2_block_ptr, (0, BLOCK_N))
+
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * Q_CTX + offs_m
+    if start_m * BLOCK_M + BLOCK_M >= Q_CTX:
+        tl.store(m_ptrs, m_i, mask=offs_m < Q_CTX)
+        tl.store(O_block_ptr, acc.to(Out.type.element_ty), boundary_check=(0,))
+    else:
+        tl.store(m_ptrs, m_i)
+        tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+def prefill_flash_forward(q1, k1, q2, k2, v, q_len, seq_len, window, sm_scale=None):
+    # shape constraints
+    Lq, Lk, Lv = q1.shape[-1], k1.shape[-1], v.shape[-1]
+    assert Lq == Lk and Lk == Lv
+    assert Lk in {16, 32, 64, 128}
+    assert q_len == seq_len or q_len == 1 
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(Lq) # the default scale factor.
+    o = torch.empty_like(q1, device=q1.device)
+    block_m = 128
+    block_n = 64 # if Lk <= 64 else 32
+    num_stages = 4 if Lk <= 64 else 3
+    num_warps = 4
+    # Tuning for H100
+    if torch.cuda.get_device_capability()[0] == 9:
+        num_warps = 8
+        num_stages = 7 if Lk >= 64 else 3
+    grid = (triton.cdiv(q1.shape[2], block_m), q1.shape[0] * q1.shape[1], 1)
+    M = torch.empty((q1.shape[0], q1.shape[1], q1.shape[2]), device=q1.device, dtype=torch.float32)
+    with torch.cuda.device(v.device.index):
+        # https://github.com/Dao-AILab/flash-attention/commit/9795159082f6e6c847db2bf4284fd17326c31fbd
+        # to avoid the device issue .
+        _attn_fwd_prefill[grid](
+            q1, k1, q2, k2, v, sm_scale, M, o,  #
+            q1.stride(0), q1.stride(1), q1.stride(2), q1.stride(3),  #
+            k1.stride(0), k1.stride(1), k1.stride(2), k1.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q1.shape[0], q1.shape[1],  #
+            Q_CTX=q_len, 
+            N_CTX=seq_len,  #
+            BLOCK_M=block_m,  #
+            BLOCK_N=block_n,  #
+            WINDOW=window,
+            BLOCK_DMODEL=Lk,  #
+            num_warps=num_warps,  #
+            num_stages=num_stages  #
+        )
+
+    return o

style.css

@@ -0,0 +1,24 @@
+h1 {
+  text-align: center;
+  display: block;
+}
+
+#duplicate-button {
+  margin: auto;
+  color: white;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+
+.contain {
+  max-width: 900px;
+  margin: auto;
+  padding-top: 1.5rem;
+}
+
+.s-pad {
+  display: block;
+  padding-top: 2rem;
+  height: 1px;
+  width: 100%;
+}