diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..1bbc33dccb83c17a9c1eb6b5fa9288111c91eebc 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,5 @@ saved_model/**/* 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 +tokenizer.json filter=lfs diff=lfs merge=lfs -text +quasar_banner.png filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..23c7f24df0ef7dc39ed320c6b1fb7585f219781f --- /dev/null +++ b/README.md @@ -0,0 +1,710 @@ +--- +language: +- en +- ar +license: mit +tags: +- silx-ai +- quasar-preview +- quasar +- foundation-model +- moe +- 18b +- 2b-active +- long-context +- bittensor +- sn24 +- decentralized-training +- distillation +- hybrid-transformer +- loop-transformer +- safe-nope +- drope +pipeline_tag: text-generation +library_name: transformers +--- + +

+ Quasar-Preview Foundation Model +

+ +# **Quasar-Preview** + +**Quasar-Preview** is the first public model in SILX AI’s **Quasar Foundation Model** series. + +It is an early preview checkpoint built to demonstrate the direction of the Quasar architecture at real scale: sparse MoE routing, hybrid recurrent/attention layers, and an experimental long-context configuration designed for future memory-based systems. + +This is **not the finished Quasar model**. + +Quasar-Preview is the first public step in a larger series of Quasar models that will continue scaling through decentralized training, distillation, architecture improvements, and long-context research on **Bittensor SN24**. + +--- + +## TL;DR + +- **First public Quasar model** +- **~18B total parameter MoE** +- **~2B active parameter path** +- **Experimental 5M-token context configuration** +- Built with **Loop Transformer + Quasar hybrid attention** +- Includes **Quasar / Raven / GLA** hybrid layers +- Designed for **Bittensor SN24 decentralized distillation** +- Trained on **>1T and <1.5T tokens** +- Long-context extension path has received **<1B tokens** so far +- Early preview checkpoint, not a final production/SOTA model + +Quasar-Preview should be understood as an **architecture preview and foundation checkpoint**, not the final endpoint of the Quasar roadmap. + +--- + +# Important Note + +Quasar-Preview is an early model from our broader Quasar model series. + +It is released to make the architecture public, allow miners and researchers to work with the model, and begin the next phase of decentralized scaling. + +This model is: + +- An **early preview checkpoint** +- The **first model** in a planned series of Quasar models +- Trained on **>1T and <1.5T tokens** +- Built for **research, distillation, and SN24 training** +- Not yet the final Quasar model +- Not intended to represent the final quality of the Quasar architecture + +Performance is expected to improve through: + +- Iterative subnet training +- Distillation cycles +- Longer training runs +- Stronger post-training +- More long-context extension training +- Future Quasar architecture updates + +--- + +# Model Overview + +| Field | Value | +| --- | --- | +| Model Name | Quasar-Preview | +| Model Family | Quasar Foundation Models | +| Organization | SILX AI | +| Model Type | `quasar_long` | +| Architecture | Quasar Long Hybrid Transformer | +| Total Parameters | ~18B class | +| Active Parameters | ~2B class sparse MoE path | +| Training Stage | Early preview checkpoint | +| Context Config | Experimental 5M-token config | +| Long-Context Method | Safe NoPE / DrOPE-style staging | +| Tokenizer | Quasar tokenizer preserved from checkpoint lineage | +| Primary Use | Research, distillation, SN24 decentralized training | +| License | MIT | + +--- + +# What Is Active In This Checkpoint? + +Quasar-Preview includes several architecture paths. Some are active in this checkpoint, while others are included for future Quasar versions. + +| Component | Status in Quasar-Preview | +| --- | --- | +| Sparse MoE | Active | +| Quasar hybrid layers | Active | +| GLA branch | Active | +| Raven branch | Active | +| GQA compatibility attention | Active in this checkpoint | +| Safe NoPE / DrOPE-style context config | Active | +| Loop Transformer scaffold | Present | +| Loop execution | Configured as single-loop | +| Looped anchor injection | Disabled | +| Engram memory | Included and loadable, not active by default | +| 5M context | Config exposed, early long-context training only | + +The goal of this release is to expose the first working Quasar architecture checkpoint while keeping the model stable for research and SN24 training. + +--- + +# Quick Start + +Quasar-Preview uses custom architecture code. + +Use `trust_remote_code=True` when loading the model. + +```python +from transformers import AutoTokenizer, AutoModelForCausalLM +import torch + +model_id = "SILX-AI/Quasar-Preview" + +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True +) + +model = AutoModelForCausalLM.from_pretrained( + model_id, + trust_remote_code=True, + torch_dtype=torch.bfloat16, + device_map="auto" +) + +prompt = "Explain the purpose of long-context models in simple terms." + +inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + +with torch.no_grad(): + output = model.generate( + **inputs, + max_new_tokens=256, + do_sample=True, + temperature=0.7, + top_p=0.9 + ) + +print(tokenizer.decode(output[0], skip_special_tokens=True)) +``` + +## Inference Notes + +Quasar-Preview is an ~18B total parameter MoE checkpoint. Even though the active path is ~2B parameters, the full checkpoint still requires loading the model weights. + +Actual memory usage depends on: + +- Precision +- Quantization +- Runtime implementation +- Sequence length +- Batch size +- Device mapping +- Whether long-context experiments are enabled + +The 5M context configuration is experimental. Do not assume ordinary inference hardware can run full 5M-token contexts without specialized infrastructure. + +--- + +# Quasar-Preview Benchmark Snapshot + +These are early benchmark results from the current Quasar checkpoint lineage. + +They should be treated as a moving snapshot, not final model quality. + +| Category | Benchmark | Quasar-Preview | +| --- | --- | ---: | +| Knowledge | MMLU (5-shot) | **68.40%** | +| Knowledge | MMLU-Pro | **33.20%** | +| Knowledge | GPQA | **25.60%** | +| Commonsense | ARC Challenge | **63.00%** | +| Commonsense | ARC Easy | **80.10%** | +| Commonsense | PIQA | **81.90%** | +| Commonsense | HellaSwag | **74.00%** | +| Science | OpenBookQA | **47.00%** | +| Math | MATH-500 (4-shot) | **71.40%** | + +## Evaluation Notes + +These results are provided as an early internal snapshot for the current Quasar-Preview checkpoint lineage. + +They are not presented as final model quality. Public verification, different harness versions, prompt formats, decoding settings, and evaluation implementations may change the reported numbers. + +When comparing Quasar-Preview to other models, please report: + +- Evaluation harness +- Harness version or commit +- Prompt format +- Shot count +- Decoding settings +- Whether chain-of-thought prompting was used +- Exact checkpoint version + +--- + +# Training Strategy + +Quasar follows a multi-stage training plan. + +Quasar-Preview is an early checkpoint from this plan. + +## Stage 1 — Base Pretraining + +The base model is trained on a broad corpus to build general next-token prediction, reasoning, and language ability. + +Goals of this stage: + +- Stabilize the sparse MoE path +- Build general language ability +- Train the hybrid Quasar stack +- Establish a checkpoint suitable for distillation and subnet training + +Quasar-Preview has been trained on **>1T and <1.5T tokens** so far. + +## Stage 2 — Distillation And Capability Training + +After base training, Quasar-Preview is improved through task distillation and targeted capability training. + +The goal is to make the checkpoint more useful for: + +- Reasoning +- Instruction-following +- Commonsense tasks +- Math and science tasks +- SN24 miner distillation +- Future post-training + +This release is designed to be a foundation for continued decentralized improvement rather than the final result. + +## Stage 3 — Long-Context Extension + +Quasar is designed to move toward ultra-long-context reasoning and memory. + +The current checkpoint exposes an experimental **5M-token context configuration** using safe NoPE / DrOPE-style staging. + +Important: the 5M context path has received **less than 1B tokens** of long-context extension training so far. + +This means the config is present, but mature 5M-token reasoning quality should not be expected yet. + +The purpose of this stage is to: + +- Preserve short-context behavior +- Avoid damaging the base model during extension +- Prepare the architecture for future long-context training +- Enable research on scalable memory and recall + +--- + +# Quasar Long Hybrid Architecture + +Quasar is a hybrid transformer architecture designed for long-context research, sparse computation, and decentralized training. + +It is built around: + +- A Loop Transformer execution scaffold +- Sparse Mixture-of-Experts routing +- Hybrid Quasar / Raven / GLA branch layers +- Optional anchor-state conditioning +- Optional Engram n-gram memory +- Safe NoPE / DrOPE-style long-context configuration + +Quasar-Preview is the first public checkpoint in this architecture family. + +--- + +# Technical Specifications + +| Component | Value | +| --- | ---: | +| Total parameters | ~18B | +| Active parameters | ~2B | +| Layers | 20 | +| Hidden size | 2048 | +| Intermediate size | 5120 | +| Attention heads | 16 | +| KV heads | 4 | +| Head dim | 128 | +| Vocabulary size | 157,184 | +| Experts | 256 | +| Experts per token | 8 | +| Shared experts | 1 | +| Active hybrid layers | 4-19 | +| Raven slots | 64 | +| Raven top-k | 32 | +| Engram slots config | 2,000,000 | +| Loop count config | 1 | +| Looped injection config | Disabled | +| Max context config | 5,000,000 | +| Safe NoPE cutoff | 512 | + +Compatibility note: this checkpoint includes GQA for the current release path. Future Quasar versions may change this component as the architecture evolves. + +--- + +# Looped Transformer Path + +Quasar includes a Loop Transformer execution path. + +The idea is to reuse the decoder stack across multiple passes, increasing effective computation depth without copying every parameter into a deeper model. + +The current checkpoint is configured conservatively: + +```text +num_loops: 1 +use_looped_injection: false +``` + +This means Quasar-Preview runs as a single-loop model by default. + +The loop machinery is still part of the architecture code and can be enabled in future Quasar configurations. + +When looped injection is enabled, Quasar keeps an anchor snapshot of the input embedding stream, usually called **P**, and injects it back into the hidden state during looped execution. + +This gives later loop passes a stable reference to the original token stream. + +The intended future looped path is: + +```text +Token IDs + | + v +Embedding Layer + | + +--> Anchor P snapshot + | + v +Decoder stack + | + v +Loop pass 1 + | + +--> inject gated Anchor P + | + v +Loop pass 2 / future passes + | + v +Final hidden state +``` + +The injection gate is initialized near zero so the model can adapt safely instead of suddenly changing behavior. + +This gives Quasar a path toward deeper effective reasoning while keeping parameter count controlled. + +--- + +# Core Data Flow + +```text +Token IDs + | + v +Token Embedding + | + +--> Optional Anchor P snapshot + | + v +Early Transformer Blocks + layers 0-3 + | + v +Hybrid Quasar Blocks + layers 4-19 + | + +--> GQA attention path + | + +--> Quasar recurrent / linear path + | + +--> Raven slot-memory path + | + +--> GLA recurrent path + | + v +Hybrid Add / Branch Merge + | + v +Optional Loop Injection / Next Loop + | + v +RMSNorm + | + v +LM Head + | + v +Next-token logits +``` + +--- + +# Hybrid Layer Composition + +The active hybrid layers are: + +```text +4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 +``` + +The current layerwise branch cycle is: + +```text +quasar -> raven -> quasar -> quasar -> gla +``` + +Across the hybrid stack, this gives: + +- **Quasar branch:** 10 layers +- **Raven branch:** 3 layers +- **GLA branch:** 3 layers + +The design keeps Quasar as the dominant branch while giving the model targeted recurrent and slot-memory paths. + +--- + +# Quasar + GLA + +GLA is used through the bundled Flash Linear Attention stack. + +The goal of the GLA branch is to give Quasar a fast recurrent sequence-mixing path that is cheaper than full dense attention at long lengths. + +Current GLA-related config: + +```text +hybrid_gla_enabled: true +hybrid_gla_expand_k: 1.0 +hybrid_gla_expand_v: 1.0 +hybrid_use_short_conv: false +``` + +GLA is not used as a standalone model here. + +It is a branch inside Quasar's hybrid layers. + +--- + +# Raven Design + +Raven is included as a slot-routed recurrent attention branch. + +Current Raven config: + +```text +hybrid_raven_enabled: true +hybrid_raven_slots: 64 +hybrid_raven_topk: 32 +hybrid_raven_decay_type: Mamba2 +``` + +Raven routes hidden states through a fixed number of recurrent memory slots. + +In this checkpoint: + +- The branch has **64 memory slots** +- It selects **top-32 routes** +- It uses a **Mamba2-style decay** + +Raven gives Quasar a memory-like path where sequence information can be compressed into routed recurrent state instead of relying only on dense attention. + +--- + +# Engram Design + +Engram is Quasar's conditional n-gram memory module. + +It is included in the repository as `engram.py` and supports: + +- n-gram orders `[2, 3]` +- 8 Engram heads +- configurable memory slots +- Triton hash-table lookup +- gated projection back into the residual stream + +Current Engram config: + +```text +engram_slots: 2,000,000 +engram_dim: 512 +engram_ngram_orders: [2, 3] +engram_num_heads: 8 +engram_residual_scale: 0.01 +engram_lr_multiplier: 5.0 +engram_layers: [] +``` + +`engram_layers` is currently empty. + +This means Engram is included and loadable, but not active by default in Quasar-Preview. + +Future Quasar versions can enable Engram on selected layers without changing the base model shape. + +Engram is intended as a fast recall path for repeated local patterns, while the main model focuses on reasoning and generalization. + +--- + +# Safe NoPE / DrOPE Context Design + +The current checkpoint uses safe NoPE as the default long-context configuration. + +Current context config: + +```text +use_nope: true +long_context_mode: rope_short_nope_long +nope_after_position: 512 +max_position_embeddings: 5,000,000 +max_seq_length: 5,000,000 +max_sequence_length: 5,000,000 +rope_scaling: null +rope_theta: 10000 +``` + +The behavior is: + +```text +Positions 0-511 + -> normal RoPE + +Positions 512+ + -> NoPE identity rotation + cos = 1 + sin = 0 +``` + +This is a safe DrOPE-style staging design for positional extension. + +The goals are: + +- Preserve short-context behavior +- Avoid stretching RoPE everywhere +- Avoid allocating a giant 5M RoPE table +- Expose a 5M sequence-length configuration +- Prepare for future long-context training runs + +Important: the 5M context path has only received **less than 1B tokens** of long-context extension training so far. + +So high-quality 5M-token reasoning should not be expected yet. + +This setting is included to expose and continue training the long-context path safely. + +--- + +# Config Snapshot + +```json +{ + "model_type": "quasar_long", + "architectures": ["QuasarLongForCausalLM"], + "hidden_size": 2048, + "intermediate_size": 5120, + "num_hidden_layers": 20, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 157184, + "num_experts": 256, + "num_experts_per_tok": 8, + "num_shared_experts": 1, + "num_loops": 1, + "use_looped_injection": false, + "hybrid_attention_layers": [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19], + "hybrid_branch_layout": "layerwise", + "hybrid_layerwise_cycle": ["quasar", "raven", "quasar", "quasar", "gla"], + "hybrid_replacement_mode": "add", + "hybrid_eval_mode": "hybrid_add", + "hybrid_quasar_enabled": true, + "hybrid_raven_enabled": true, + "hybrid_gla_enabled": true, + "hybrid_raven_slots": 64, + "hybrid_raven_topk": 32, + "use_nope": true, + "long_context_mode": "rope_short_nope_long", + "nope_after_position": 512, + "max_position_embeddings": 5000000, + "max_seq_length": 5000000, + "max_sequence_length": 5000000 +} +``` + +--- + +# Intended Use + +Quasar-Preview is designed as an early foundation checkpoint for the Quasar ecosystem. + +It is primarily intended for: + +- **Bittensor SN24 miners** participating in decentralized training and knowledge distillation +- **Distillation pipelines** transferring capabilities from stronger teacher models +- **Research on long-context architectures** +- **Research on sparse MoE systems** +- **Hybrid attention research** +- **Agentic system experiments** +- **Memory and recall experiments** +- **Future Quasar model development** + +This model is best treated as a research and development checkpoint. + +--- + +# Out-of-Scope Use + +Quasar-Preview is not intended to be used as: + +- A final production assistant +- A safety-aligned chatbot +- A medical, legal, or financial authority +- A final benchmark-maximized release +- Proof of mature 5M-token reasoning quality +- The final Quasar architecture endpoint + +The model may produce incorrect, unsafe, biased, or low-quality outputs. + +Use appropriate evaluation, filtering, and safety layers before any deployment. + +--- + +# Limitations + +Quasar-Preview is early. + +Known limitations: + +- It is not the finished Quasar model. +- It is the first model in a broader Quasar series. +- Long-context behavior is experimental. +- The 5M-token context is a configuration path, not yet mature 5M-token reasoning quality. +- The long-context path has received less than 1B tokens of extension training so far. +- Some architecture modules are included for future versions but disabled in this checkpoint. +- Engram is included but not active by default. +- Loop execution is configured as single-loop by default. +- Benchmarks are early checkpoint-lineage snapshots and require public verification. +- The model may hallucinate or produce incorrect answers. +- The model has not completed the full Quasar training roadmap. + +--- + +# Bittensor SN24 + +Quasar-Preview is designed for the **SN24 Quasar subnet** on Bittensor. + +The goal is to create a shared architecture where miners can continuously improve the model through distributed knowledge distillation, evaluation, and iterative training. + +SN24 is intended to support: + +- Open model improvement +- Competitive distillation +- Decentralized training incentives +- Shared progress on the Quasar architecture +- Long-context and memory-focused model development + +Quasar-Preview is the starting checkpoint for this direction. + +--- + +# Roadmap + +Quasar-Preview is only the first public model in the Quasar series. + +Next Quasar models will continue toward: + +- Larger-scale decentralized training +- More training tokens +- Stronger post-training +- Better reasoning performance +- More stable long-context behavior +- More long-context extension training +- Deeper Loop Transformer experiments +- More Raven, GLA, and Engram experimentation +- Improved benchmark performance +- Stronger agentic and memory capabilities + +Future releases may change architecture components, routing, loop configuration, long-context training strategy, and active memory modules as the Quasar series evolves. + +--- + +# Release Statement + +Quasar-Preview is not the final destination. + +It is the first public checkpoint in the Quasar model series and the first public proof of the architecture direction at scale. + +The model is early, but it is real, usable, and ready for research, distillation, and decentralized improvement. + +This is the beginning of Quasar. diff --git a/chat_template.jinja b/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..71f4509d87bdac6c327e984efd20a934c42139de --- /dev/null +++ b/chat_template.jinja @@ -0,0 +1 @@ +{% for message in messages %}{% set role = message['role'] | lower %}{% if role == 'user' %}{% set role = 'HUMAN' %}{% endif %}{% set role = role | upper %}{{ '' + role + '' + message['content'] }}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT' }}{% endif %} \ No newline at end of file diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..449d5a64b757c94f195cf3481302e6be7fd20c2a --- /dev/null +++ b/config.json @@ -0,0 +1,116 @@ +{ + "_name_or_path": "silx-ai/Quasar-T", + "architectures": [ + "QuasarLongForCausalLM" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_quasar_long.QuasarLongConfig", + "AutoModelForCausalLM": "modeling_quasar_long.QuasarLongForCausalLM" + }, + "dtype": "bfloat16", + "embedding_dropout": 0.0, + "engram_dim": 512, + "engram_layers": [], + "engram_lr_multiplier": 5.0, + "engram_ngram_orders": [ + 2, + 3 + ], + "engram_num_heads": 8, + "engram_residual_scale": 0.01, + "engram_slots": 2000000, + "eos_token_id": 156892, + "first_k_dense_replace": 1, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2048, + "hybrid_alpha_init": 0.8473, + "hybrid_attention_layers": [ + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19 + ], + "hybrid_branch_layout": "layerwise", + "hybrid_eval_force_branch": "", + "hybrid_eval_mode": "hybrid_add", + "hybrid_gla_enabled": true, + "hybrid_gla_expand_k": 1.0, + "hybrid_gla_expand_v": 1.0, + "hybrid_layerwise_cycle": [ + "quasar", + "raven", + "quasar", + "quasar", + "gla" + ], + "hybrid_quasar_enabled": true, + "hybrid_raven_decay_type": "Mamba2", + "hybrid_raven_enabled": true, + "hybrid_raven_slots": 64, + "hybrid_raven_topk": 32, + "hybrid_replacement_mode": "add", + "hybrid_use_short_conv": false, + "initializer_range": 0.02, + "intermediate_size": 5120, + "long_context_mode": "rope_short_nope_long", + "max_position_embeddings": 5000000, + "max_seq_length": 5000000, + "max_sequence_length": 5000000, + "max_window_layers": 20, + "model_type": "quasar_long", + "moe_intermediate_size": 512, + "moe_router_enable_expert_bias": true, + "moe_shared_expert_intermediate_size": 512, + "mtp_loss_scaling_factor": 0, + "n_group": 8, + "nope_after_position": 512, + "norm_topk_prob": true, + "num_attention_heads": 16, + "num_experts": 256, + "num_experts_per_tok": 8, + "num_hidden_layers": 20, + "num_key_value_heads": 4, + "num_loops": 1, + "num_nextn_predict_layers": 0, + "num_shared_experts": 1, + "output_dropout": 0.0, + "output_router_logits": false, + "pad_token_id": 156892, + "partial_rotary_factor": 0.5, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "partial_rotary_factor": 0.5, + "rope_theta": 10000, + "rope_type": "default" + }, + "rope_scaling": null, + "rope_theta": 10000, + "routed_scaling_factor": 2.5, + "router_dtype": "fp32", + "score_function": "sigmoid", + "tie_word_embeddings": false, + "topk_group": 4, + "transformers_version": "5.10.2", + "use_bias": false, + "use_cache": true, + "use_looped_injection": false, + "use_nope": true, + "use_qk_norm": true, + "use_qkv_bias": false, + "use_rmsnorm": true, + "vocab_size": 157184 +} diff --git a/configuration_quasar_long.py b/configuration_quasar_long.py new file mode 100644 index 0000000000000000000000000000000000000000..e64226417fbb5f6bce3397d1b42213dca74b81c3 --- /dev/null +++ b/configuration_quasar_long.py @@ -0,0 +1,140 @@ +"""Quasar Long model configuration""" + +from transformers.configuration_utils import PretrainedConfig + + +class QuasarLongConfig(PretrainedConfig): + model_type = "quasar_long" + + def __init__( + self, + vocab_size=157184, + hidden_size=2048, + intermediate_size=5120, + num_hidden_layers=20, + num_attention_heads=16, + num_key_value_heads=4, + hidden_act="silu", + use_qkv_bias=False, # quasar legacy + use_bias=False, # quasar legacy + rms_norm_eps=1e-06, + tie_word_embeddings=False, # PretrainedConfig key, here change default value. + embedding_dropout=0.0, + attention_dropout=0.0, + output_dropout=0.0, + initializer_range=0.02, + max_position_embeddings=32768, + rope_theta=600000.0, + use_cache=True, + max_window_layers=20, + rope_scaling=None, + pad_token_id=156892, + eos_token_id=156892, + num_experts=256, + num_shared_experts=1, + num_experts_per_tok=8, + n_group=8, + topk_group=4, + moe_intermediate_size=512, + first_k_dense_replace=1, + head_dim=128, + output_router_logits=False, + use_qk_norm=True, + num_nextn_predict_layers=0, + mtp_loss_scaling_factor=0, + moe_router_enable_expert_bias=True, + routed_scaling_factor=1.0, + hybrid_attention_layers=None, + hybrid_alpha_init=-15.0, + hybrid_gla_expand_k=1.0, + hybrid_gla_expand_v=1.0, + hybrid_use_short_conv=False, + hybrid_quasar_enabled=True, + hybrid_gla_enabled=True, + hybrid_branch_layout="mixed", + hybrid_layerwise_cycle=None, + # ── Looped Transformer ──────────────────────────────────────────────── + num_loops=1, + use_looped_injection=False, + # ── Engram Conditional Memory ───────────────────────────────────────── + # engram_layers=[] → module disabled (zero overhead, backward-compatible). + engram_layers=None, + engram_dim=512, + engram_slots=2_000_000, + engram_num_heads=8, + engram_ngram_orders=None, + engram_lr_multiplier=5.0, + use_nope=False, + long_context_mode="rope_short_nope_long", + nope_after_position=512, + max_seq_length=None, + max_sequence_length=None, + **kwargs, + ): + self.num_hidden_layers = num_hidden_layers + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.use_qkv_bias = use_qkv_bias + self.use_bias = use_bias + self.rms_norm_eps = rms_norm_eps + self.embedding_dropout = embedding_dropout + self.attention_dropout = attention_dropout + self.output_dropout = output_dropout + self.num_nextn_predict_layers = num_nextn_predict_layers + self.mtp_loss_scaling_factor = mtp_loss_scaling_factor + self.initializer_range = initializer_range + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.use_cache = use_cache + self.max_window_layers = max_window_layers + self.head_dim = head_dim or self.hidden_size // self.num_attention_heads + self.rope_scaling = rope_scaling + self.use_qk_norm = use_qk_norm + self.moe_router_enable_expert_bias = moe_router_enable_expert_bias + self.routed_scaling_factor = routed_scaling_factor + self.hybrid_attention_layers = hybrid_attention_layers or [] + self.hybrid_alpha_init = hybrid_alpha_init + self.hybrid_gla_expand_k = hybrid_gla_expand_k + self.hybrid_gla_expand_v = hybrid_gla_expand_v + self.hybrid_use_short_conv = hybrid_use_short_conv + self.hybrid_quasar_enabled = hybrid_quasar_enabled + self.hybrid_gla_enabled = hybrid_gla_enabled + self.hybrid_branch_layout = hybrid_branch_layout + self.hybrid_layerwise_cycle = list(hybrid_layerwise_cycle) if hybrid_layerwise_cycle is not None else [ + "quasar", + "raven", + "gla", + ] + + # Looped Transformer + self.num_loops = num_loops + self.use_looped_injection = use_looped_injection + + # Engram Conditional Memory + self.engram_layers = list(engram_layers) if engram_layers is not None else [] + self.engram_dim = engram_dim + self.engram_slots = engram_slots + self.engram_num_heads = engram_num_heads + self.engram_ngram_orders = list(engram_ngram_orders) if engram_ngram_orders is not None else [2, 3] + self.engram_lr_multiplier = engram_lr_multiplier + self.use_nope = use_nope + self.long_context_mode = long_context_mode + self.nope_after_position = int(nope_after_position) + self.max_seq_length = int(max_seq_length) if max_seq_length is not None else None + self.max_sequence_length = int(max_sequence_length) if max_sequence_length is not None else None + + # MoE configs + self.num_experts = num_experts + self.num_shared_experts = num_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.n_group = n_group + self.topk_group = topk_group + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.output_router_logits = output_router_logits + + super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/engram.py b/engram.py new file mode 100644 index 0000000000000000000000000000000000000000..a8e3316b51fed495793e67feca040858f816509f --- /dev/null +++ b/engram.py @@ -0,0 +1,611 @@ +""" +EngramModule: Conditional N-gram Memory for Quasar-RoPE +Implements Engram from DeepSeek-AI (arXiv:2601.07372). + +Design constraints: + - No Python loops over T (sequence length) or B (batch). + - N-gram extraction via torch.unfold (single vectorized op). + - Hash computed via vectorized XOR reduction (loop over n=2..3 only, compile-time constant). + - Embedding lookup via batched advanced indexing — no loop over T. + - Optional Triton kernel fuses hash + lookup + accumulation into a single SRAM pass. + - Zero output at init: conv.weight=0, out_proj uses deep Trinity init. +""" + +import math +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _next_prime(n: int) -> int: + """Smallest prime >= n.""" + def _is_prime(x: int) -> bool: + if x < 2: + return False + if x == 2: + return True + if x % 2 == 0: + return False + for i in range(3, int(x ** 0.5) + 1, 2): + if x % i == 0: + return False + return True + n = max(n, 2) + while not _is_prime(n): + n += 1 + return n + + +# ───────────────────────────────────────────────────────────────────────────── +# Triton Kernel: Fused N-gram Hash + Embedding Lookup +# +# Grid: (B, T). Each program handles one (batch, position) pair. +# For each of the `num_tables` embedding tables: +# 1. Load the suffix n-gram ending at position t (causal, no future tokens). +# 2. Compute XOR-multiplicative hash (loop over n ≤ 3, constexpr-unrolled). +# 3. Index into the embedding table and write directly to output. +# One SRAM pass — no intermediate [B,T,n] tensor, no round-trip to HBM. +# ───────────────────────────────────────────────────────────────────────────── + +if HAS_TRITON: + @triton.jit + def _engram_hash_lookup_kernel( + # [B, T] canonical token IDs (int32 on device) + canonical_ptr, stride_cb, stride_ct, + # [B, T, num_tables * d_slot] output (bfloat16) + output_ptr, stride_ob, stride_ot, + # [num_tables, M, d_slot] embedding tables (float32) + tables_ptr, stride_tn, stride_tm, stride_td, + # [num_tables] per-table seeds (int64) + seeds_ptr, + # [num_ngram_orders] ngram order values, e.g. [2, 3] + ngrams_ptr, + # Scalars + B, T: tl.constexpr, M, d_slot: tl.constexpr, + num_tables: tl.constexpr, + num_ngram_orders: tl.constexpr, # ≤ 4 + num_heads: tl.constexpr, # ≤ 16 + MAX_N: tl.constexpr, # max(ngram_orders), e.g. 3 + BLOCK_D: tl.constexpr, # power-of-2 ≥ d_slot + ): + b_idx = tl.program_id(0) + t_idx = tl.program_id(1) + + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < d_slot + + # Pre-load the last MAX_N canonical tokens ending at t_idx (causal). + # Positions before 0 are treated as padding (0). We unroll to pure scalars for Triton compatibility. + # Crucial Safety: clamp pos using tl.where to ensure pointer arithmetic is never negative. + c0 = tl.full((), 0, dtype=tl.int64) + c1 = tl.full((), 0, dtype=tl.int64) + c2 = tl.full((), 0, dtype=tl.int64) + c3 = tl.full((), 0, dtype=tl.int64) + + if MAX_N >= 1: + pos_raw = t_idx - (MAX_N - 1 - 0) + valid = pos_raw >= 0 + pos = tl.where(valid, pos_raw, 0) + tok = tl.load( + canonical_ptr + b_idx * stride_cb + pos * stride_ct, + mask=valid, other=0, + ) + c0 = tl.where(valid, tok.to(tl.int64), tl.full((), 0, dtype=tl.int64)) + if MAX_N >= 2: + pos_raw = t_idx - (MAX_N - 1 - 1) + valid = pos_raw >= 0 + pos = tl.where(valid, pos_raw, 0) + tok = tl.load( + canonical_ptr + b_idx * stride_cb + pos * stride_ct, + mask=valid, other=0, + ) + c1 = tl.where(valid, tok.to(tl.int64), tl.full((), 0, dtype=tl.int64)) + if MAX_N >= 3: + pos_raw = t_idx - (MAX_N - 1 - 2) + valid = pos_raw >= 0 + pos = tl.where(valid, pos_raw, 0) + tok = tl.load( + canonical_ptr + b_idx * stride_cb + pos * stride_ct, + mask=valid, other=0, + ) + c2 = tl.where(valid, tok.to(tl.int64), tl.full((), 0, dtype=tl.int64)) + if MAX_N >= 4: + pos_raw = t_idx - (MAX_N - 1 - 3) + valid = pos_raw >= 0 + pos = tl.where(valid, pos_raw, 0) + tok = tl.load( + canonical_ptr + b_idx * stride_cb + pos * stride_ct, + mask=valid, other=0, + ) + c3 = tl.where(valid, tok.to(tl.int64), tl.full((), 0, dtype=tl.int64)) + + # Iterate over all tables; loop bounds are constexpr → fully unrolled by compiler. + for n_ord in tl.static_range(4): # ≤ num_ngram_orders + if n_ord < num_ngram_orders: + n = tl.load(ngrams_ptr + n_ord).to(tl.int32) + + for k in tl.static_range(16): # ≤ num_heads + if k < num_heads: + # Safety: Compute unique table_idx directly from static loop indices to avoid mutable variable register compilation bugs. + table_idx = n_ord * num_heads + k + seed = tl.load(seeds_ptr + table_idx).to(tl.int64) + + # XOR-multiplicative hash over the suffix n-gram. + # loop over MAX_N positions; positions outside the suffix are skipped. + h = seed + for i in tl.static_range(MAX_N): + include = i >= (MAX_N - n) + tok = tl.full((), 0, dtype=tl.int64) + if i == 0: + tok = c0 + elif i == 1: + tok = c1 + elif i == 2: + tok = c2 + elif i == 3: + tok = c3 + new_h = h * 2654435761 ^ tok + h = tl.where(include, new_h, h) + + # Clamp absolute value of hash using tl.where for maximum Triton version safety + idx = tl.where(h >= 0, h, -h) % M + + # Load d_slot floats from embed_tables[table_idx, idx] + emb_base = table_idx * stride_tn + idx * stride_tm + emb = tl.load( + tables_ptr + emb_base + d_offs * stride_td, + mask=d_mask, other=0.0, + ) + + # Write to output[b, t, table_idx*d_slot : (table_idx+1)*d_slot] + out_base = b_idx * stride_ob + t_idx * stride_ot + table_idx * d_slot + tl.store( + output_ptr + out_base + d_offs, + emb.to(output_ptr.dtype.element_ty), + mask=d_mask, + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Custom Autograd Function for Fused Triton Training Lookup +# ───────────────────────────────────────────────────────────────────────────── + +class FusedEngramLookupFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + canonical, + embed_tables, + seeds, + ngram_orders_buf, + M, + d_slot, + num_tables, + num_ngram_orders, + num_heads, + ngram_orders, + ): + ctx.save_for_backward(canonical, seeds, ngram_orders_buf) + ctx.M = M + ctx.d_slot = d_slot + ctx.num_tables = num_tables + ctx.num_ngram_orders = num_ngram_orders + ctx.num_heads = num_heads + ctx.ngram_orders = ngram_orders + ctx.embed_tables_shape = embed_tables.shape + + B, T = canonical.shape + BLOCK_D = triton.next_power_of_2(d_slot) + + out = torch.empty( + B, T, num_tables * d_slot, + device=canonical.device, dtype=embed_tables.dtype, + ) + + tables = embed_tables.contiguous() + + _engram_hash_lookup_kernel[(B, T)]( + canonical.int().contiguous(), canonical.stride(0), canonical.stride(1), + out, out.stride(0), out.stride(1), + tables, tables.stride(0), tables.stride(1), tables.stride(2), + seeds.contiguous(), + ngram_orders_buf.contiguous(), + B, T, M, d_slot, + num_tables, num_ngram_orders, num_heads, + MAX_N=max(ngram_orders), + BLOCK_D=BLOCK_D, + ) + return out + + @staticmethod + def backward(ctx, grad_output): + canonical, seeds, ngram_orders_buf = ctx.saved_tensors + B, T = canonical.shape + device = canonical.device + + # 1. Re-compute hashes for each table in vectorized form + all_hashes = torch.empty(ctx.num_tables, B * T, dtype=torch.long, device=device) + table_idx = 0 + for n_idx, n in enumerate(ctx.ngram_orders): + padded = F.pad(canonical, (n - 1, 0), value=0) + ngrams = padded.unfold(dimension=1, size=n, step=1) + for k in range(ctx.num_heads): + seed = int(seeds[table_idx].item()) + h = torch.full(ngrams.shape[:2], seed, dtype=torch.long, device=device) + for i in range(ngrams.shape[-1]): + h = h * 2654435761 ^ ngrams[..., i] + h = h.abs() % ctx.M + all_hashes[table_idx] = h.view(B * T) + table_idx += 1 + + # 2. Reshape and permute grad_output from [B, T, num_tables * d_slot] back to [num_tables, B * T, d_slot] + grad_out_reshaped = grad_output.reshape(B, T, ctx.num_tables, ctx.d_slot).permute(2, 0, 1, 3).reshape(ctx.num_tables, B * T, ctx.d_slot) + + # 3. Accumulate gradients into grad_embed_tables using PyTorch's native CUDA-optimized index_put_ scatter-add + grad_embed_tables = torch.zeros(ctx.embed_tables_shape, dtype=grad_output.dtype, device=device) + tbl_idx = torch.arange(ctx.num_tables, device=device).unsqueeze(1).expand(ctx.num_tables, B * T) + + grad_embed_tables.index_put_((tbl_idx, all_hashes), grad_out_reshaped, accumulate=True) + + # Return gradients matching forward arguments (None for non-tensor / constant arguments) + return None, grad_embed_tables, None, None, None, None, None, None, None, None + + +# ───────────────────────────────────────────────────────────────────────────── +# Lightweight RMSNorm (standalone; avoids circular import from quasar_rope) +# ───────────────────────────────────────────────────────────────────────────── + +class _RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +# ───────────────────────────────────────────────────────────────────────────── +# EngramModule +# ───────────────────────────────────────────────────────────────────────────── + +class EngramModule(nn.Module): + """ + Engram Conditional Memory Module (DeepSeek-AI, arXiv:2601.07372). + + Replaces expensive attention layers for static N-gram patterns with + O(1) hash-table lookups gated into the hidden state. + + All operations are fully vectorized — no Python loops over T or B: + • N-gram extraction: torch.unfold (single op) + • Hash computation: vectorized XOR accumulation (loop over n=2..3 only) + • Embedding lookup: batched advanced indexing (single gather) + • Conv: nn.Conv1d with causal pad + slice + """ + + def __init__( + self, + vocab_size: int, + d_model: int, + d_mem: int, + num_heads: int = 8, + ngram_orders: list = None, + target_slots: int = 5_700_000, + n_layers: int = 24, + ): + super().__init__() + + if ngram_orders is None: + ngram_orders = [2, 3] + + self.vocab_size = vocab_size + self.d_model = d_model + self.d_mem = d_mem + self.num_heads = num_heads + self.ngram_orders = list(ngram_orders) + self.num_ngram_orders = len(ngram_orders) + self.num_tables = self.num_ngram_orders * num_heads + self.n_layers = n_layers + + # ── A. Tokenizer Compression Buffer ────────────────────────────────── + # Surjective P: V → V', ~23% compression. + # Deterministic multiplicative hash — no tokenizer object needed at + # construction time (avoids FSDP serialization problems). + compressed_size = max(1, int(vocab_size * 0.77)) + self.compressed_vocab_size = compressed_size + token_map = ( + torch.arange(vocab_size, dtype=torch.long) * 2654435761 + ) % compressed_size + self.register_buffer('token_map', token_map) + + # ── B. Embedding Tables ─────────────────────────────────────────────── + # All num_tables share the same prime size M for vectorized indexing. + slots_per_table = max(1, target_slots // self.num_tables) + self.M = _next_prime(slots_per_table) + self.d_slot = max(16, d_mem // max(1, self.num_tables)) + self.total_embed_dim = self.num_tables * self.d_slot + + # Single parameter tensor — enables batched advanced-index gather. + self.embed_tables = nn.Parameter( + torch.empty(self.num_tables, self.M, self.d_slot) + ) + + # Per-table hashing seeds (non-trainable). + seeds = torch.randint(1, 2 ** 31 - 1, (self.num_tables,), dtype=torch.long) + self.register_buffer('seeds', seeds) + + # N-gram order list as a buffer for the Triton kernel. + self.register_buffer( + 'ngram_orders_buf', + torch.tensor(self.ngram_orders, dtype=torch.long), + ) + + # ── C. Projection total_embed_dim → d_mem ──────────────────────────── + self.embed_proj = nn.Linear(self.total_embed_dim, d_mem, bias=False) + + # ── D. Context-aware gating ─────────────────────────────────────────── + self.q_proj = nn.Linear(d_model, d_mem, bias=False) + self.W_K = nn.Linear(d_mem, d_mem, bias=False) + self.W_V = nn.Linear(d_mem, d_mem, bias=False) + + # ── E. Causal depthwise Conv1d ──────────────────────────────────────── + # kernel=4, dilation=3 → causal receptive field = 1 + (4-1)*3 = 10 + self.kernel_size = 4 + self.dilation = 3 + self.conv_norm = _RMSNorm(d_mem) + self.conv = nn.Conv1d( + d_mem, d_mem, + kernel_size=self.kernel_size, + dilation=self.dilation, + groups=d_mem, # depthwise + bias=False, + ) + + # ── F. Output projection d_mem → d_model ───────────────────────────── + self.out_proj = nn.Linear(d_mem, d_model, bias=False) + + # Triton eligible when compile-time bounds fit the kernel + self._triton_ok = ( + HAS_TRITON + and self.num_ngram_orders <= 4 + and num_heads <= 16 + and max(ngram_orders) <= 3 + ) + self.triton_training = True + + self._init_weights() + + # ── Initialization ──────────────────────────────────────────────────────── + + def _init_weights(self): + # 1. Deterministic buffer re-population (bypasses meta-device empty uninitialized memory) + if hasattr(self, "token_map") and self.token_map is not None: + dev = "cpu" if self.token_map.device.type == "meta" else self.token_map.device + t_map = (torch.arange(self.vocab_size, dtype=torch.long, device=dev) * 2654435761) % self.compressed_vocab_size + self.token_map.data.copy_(t_map) + + if hasattr(self, "seeds") and self.seeds is not None: + # Deterministic hash seeds across all ranks + g = torch.Generator().manual_seed(42) + dev = "cpu" if self.seeds.device.type == "meta" else self.seeds.device + s_t = torch.randint(1, 2 ** 31 - 1, (self.num_tables,), dtype=torch.long, device=dev, generator=g) + self.seeds.data.copy_(s_t) + + if hasattr(self, "ngram_orders_buf") and self.ngram_orders_buf is not None: + dev = "cpu" if self.ngram_orders_buf.device.type == "meta" else self.ngram_orders_buf.device + ord_buf = torch.tensor(self.ngram_orders, dtype=torch.long, device=dev) + self.ngram_orders_buf.data.copy_(ord_buf) + + trinity_std = 0.5 / math.sqrt(self.d_model) + scale_factor = 1.0 / math.sqrt(2 * self.n_layers) + + # 2. Deep init on output → zero-init to guarantee exactly zero output at step 0 + nn.init.zeros_(self.out_proj.weight) + # Gating projections: standard Trinity + nn.init.normal_(self.q_proj.weight, std=trinity_std) + nn.init.normal_(self.W_K.weight, std=trinity_std) + nn.init.normal_(self.W_V.weight, std=trinity_std) + # embed_proj: standard Trinity + nn.init.normal_(self.embed_proj.weight, std=trinity_std) + # Conv: zero init → identity pass-through at step 0 + nn.init.zeros_(self.conv.weight) + # Embedding tables: small normal (paper standard) + nn.init.normal_(self.embed_tables, std=0.01) + # Conv norm: fill with ones + if hasattr(self.conv_norm, "weight") and self.conv_norm.weight is not None: + nn.init.ones_(self.conv_norm.weight) + + # Check for any non-finite initialization values + for name, p in [("out_proj", self.out_proj.weight), ("q_proj", self.q_proj.weight), + ("W_K", self.W_K.weight), ("W_V", self.W_V.weight), + ("embed_proj", self.embed_proj.weight), ("conv", self.conv.weight), + ("embed_tables", self.embed_tables), ("conv_norm", self.conv_norm.weight)]: + if p.device.type != "meta": + if not torch.isfinite(p).all(): + print(f"[engram-init-warn] Parameter {name} contains non-finite values! Re-initializing with zeros.", flush=True) + nn.init.zeros_(p) + + # ── Core helpers ────────────────────────────────────────────────────────── + + @staticmethod + def _hash_ngrams(ngrams: torch.Tensor, table_size: int, seed: int) -> torch.Tensor: + """ + Vectorized XOR-multiplicative hash. + ngrams: [B, T, n] — n ∈ {2, 3}, compile-time constant. + Returns: [B, T] — indices into embedding table. + No loop over T or B; only loops over n (≤ 3). + """ + h = torch.full(ngrams.shape[:2], seed, dtype=torch.long, device=ngrams.device) + for i in range(ngrams.shape[-1]): # n iterations, NOT T + h = h * 2654435761 ^ ngrams[..., i] + return h.abs() % table_size + + def _lookup_pytorch(self, canonical: torch.Tensor) -> torch.Tensor: + """ + Pure-PyTorch path: fully vectorized, no T/B loops. + + Steps: + 1. For each n-gram order, extract suffix n-grams via unfold → [B, T, n] + 2. Hash all (n, k) pairs → [num_tables, B, T] + 3. Batched advanced-index gather from embed_tables → [num_tables, B*T, d_slot] + 4. Reshape to [B, T, total_embed_dim] + """ + B, T = canonical.shape + device = canonical.device + + # Step 1+2: collect hashes for all tables — loop over num_tables (≤ 32, not over T) + all_hashes = torch.empty(self.num_tables, B * T, dtype=torch.long, device=device) + table_idx = 0 + seeds_cpu = self.seeds.cpu().tolist() if hasattr(self, "seeds") and self.seeds is not None else [] + for n_idx, n in enumerate(self.ngram_orders): # 2 or 3 iterations + # Vectorized n-gram extraction: unfold over T → [B, T, n] + padded = F.pad(canonical, (n - 1, 0), value=0) # [B, T+n-1] + ngrams = padded.unfold(dimension=1, size=n, step=1) # [B, T, n] + + for k in range(self.num_heads): # num_heads iterations (≤ 16) + seed = seeds_cpu[table_idx] if table_idx < len(seeds_cpu) else 42 + h = self._hash_ngrams(ngrams, self.M, seed) # [B, T] + all_hashes[table_idx] = h.view(B * T) + table_idx += 1 + + # Step 3: Single batched gather — no loop over T + # embed_tables: [num_tables, M, d_slot] + # all_hashes: [num_tables, B*T] + # Expand table index for advanced indexing + tbl_idx = torch.arange(self.num_tables, device=device).unsqueeze(1).expand( + self.num_tables, B * T + ) # [num_tables, B*T] + embeddings = self.embed_tables[tbl_idx, all_hashes] # [num_tables, B*T, d_slot] + + # Step 4: Reshape to [B, T, total_embed_dim] + embeddings = embeddings.permute(1, 0, 2) # [B*T, num_tables, d_slot] + return embeddings.reshape(B, T, self.total_embed_dim) + + def _lookup_triton(self, canonical: torch.Tensor) -> torch.Tensor: + """ + Triton path: fused hash + lookup in a single SRAM pass. + Uses FusedEngramLookupFunction to support exact backward auto-differentiation in training. + """ + return FusedEngramLookupFunction.apply( + canonical, + self.embed_tables, + self.seeds, + self.ngram_orders_buf, + self.M, + self.d_slot, + self.num_tables, + self.num_ngram_orders, + self.num_heads, + self.ngram_orders, + ) + + # ── Forward ─────────────────────────────────────────────────────────────── + + def forward( + self, + input_ids: torch.Tensor, # [B, T] raw token IDs + hidden_states: torch.Tensor, # [B, T, d_model] + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Returns: + engram_out : [B, T, d_model] — to add to residual stream + alpha_mean : scalar tensor — mean gate value for LatentMemory suppression + """ + B, T = input_ids.shape + orig_dtype = hidden_states.dtype + + # ── A. Token compression ───────────────────────────────────────────── + # Single gather op — no loop + canonical = self.token_map[input_ids.clamp(0, self.vocab_size - 1)] # [B, T] + + # ── B+C. Hash → lookup → project to d_mem ─────────────────────────── + use_triton_lookup = self._triton_ok and canonical.is_cuda and ( + self.training is False or bool(getattr(self, "triton_training", False)) + ) + if use_triton_lookup: + raw_embed = self._lookup_triton(canonical) # [B, T, total_embed_dim] + else: + raw_embed = self._lookup_pytorch(canonical) # [B, T, total_embed_dim] + + raw_embed = raw_embed.to(orig_dtype) + + debug_engram = bool(int(os.environ.get("ENGRAM_DEBUG", "0"))) + if debug_engram: + print(f"[engram-debug] embed_tables: finite={torch.isfinite(self.embed_tables).all().item()} min={self.embed_tables.float().min().item():.6g} max={self.embed_tables.float().max().item():.6g}", flush=True) + print(f"[engram-debug] hidden_states: finite={torch.isfinite(hidden_states).all().item()} min={hidden_states.float().min().item():.6g} max={hidden_states.float().max().item():.6g}", flush=True) + print(f"[engram-debug] raw_embed: finite={torch.isfinite(raw_embed).all().item()} min={raw_embed.float().min().item():.6g} max={raw_embed.float().max().item():.6g}", flush=True) + + raw_embed = torch.nan_to_num(raw_embed, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-10.0, 10.0) + e_t = self.embed_proj(raw_embed) # [B, T, d_mem] + if debug_engram: + print(f"[engram-debug] e_t: finite={torch.isfinite(e_t).all().item()} min={e_t.float().min().item():.6g} max={e_t.float().max().item():.6g}", flush=True) + e_t = torch.nan_to_num(e_t, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + + # ── D. Context-aware gating ────────────────────────────────────────── + h_proj = self.q_proj(hidden_states) # [B, T, d_mem] + k_t = self.W_K(e_t) # [B, T, d_mem] + v_t = self.W_V(e_t) # [B, T, d_mem] + if debug_engram: + print(f"[engram-debug] h_proj: finite={torch.isfinite(h_proj).all().item()} min={h_proj.float().min().item():.6g} max={h_proj.float().max().item():.6g}", flush=True) + print(f"[engram-debug] k_t: finite={torch.isfinite(k_t).all().item()} min={k_t.float().min().item():.6g} max={k_t.float().max().item():.6g}", flush=True) + print(f"[engram-debug] v_t: finite={torch.isfinite(v_t).all().item()} min={v_t.float().min().item():.6g} max={v_t.float().max().item():.6g}", flush=True) + + h_proj = torch.nan_to_num(h_proj, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + k_t = torch.nan_to_num(k_t, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + v_t = torch.nan_to_num(v_t, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + + # L2-normalize for stability (matches Quasar key normalization) + q_norm = F.normalize(h_proj.float(), dim=-1, eps=1e-6).to(orig_dtype) + k_norm = F.normalize(k_t.float(), dim=-1, eps=1e-6).to(orig_dtype) + + # Scalar gate per token per position + alpha_logits = (q_norm * k_norm).sum(-1, keepdim=True).float() / math.sqrt(self.d_mem) + alpha_t = torch.sigmoid(alpha_logits.clamp_(-30.0, 30.0)).to(orig_dtype) # [B, T, 1] + if debug_engram: + print(f"[engram-debug] alpha_t: finite={torch.isfinite(alpha_t).all().item()} min={alpha_t.float().min().item():.6g} max={alpha_t.float().max().item():.6g}", flush=True) + gated = alpha_t * v_t # [B, T, d_mem] + gated = torch.nan_to_num(gated, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + + # ── E. Causal depthwise conv ───────────────────────────────────────── + # Fully vectorized: F.pad + Conv1d + slice — no loop over T + causal_pad = (self.kernel_size - 1) * self.dilation + g_norm = self.conv_norm(gated) # [B, T, d_mem] + if debug_engram: + print(f"[engram-debug] gated: finite={torch.isfinite(gated).all().item()} min={gated.float().min().item():.6g} max={gated.float().max().item():.6g}", flush=True) + print(f"[engram-debug] conv_norm.weight: finite={torch.isfinite(self.conv_norm.weight).all().item()} min={self.conv_norm.weight.float().min().item():.6g} max={self.conv_norm.weight.float().max().item():.6g}", flush=True) + print(f"[engram-debug] g_norm: finite={torch.isfinite(g_norm).all().item()} min={g_norm.float().min().item():.6g} max={g_norm.float().max().item():.6g}", flush=True) + g_norm = torch.nan_to_num(g_norm, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + g_t = g_norm.transpose(1, 2) # [B, d_mem, T] + g_t = F.pad(g_t, (causal_pad, 0)) # [B, d_mem, T+pad] + g_t = self.conv(g_t)[..., :T] # [B, d_mem, T] + g_t = F.silu(g_t).transpose(1, 2) # [B, T, d_mem] + Y = g_t + gated # residual + if debug_engram: + print(f"[engram-debug] Y: finite={torch.isfinite(Y).all().item()} min={Y.float().min().item():.6g} max={Y.float().max().item():.6g}", flush=True) + Y = torch.nan_to_num(Y, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + + # ── F. Output projection ───────────────────────────────────────────── + engram_out = self.out_proj(Y) # [B, T, d_model] + if debug_engram: + print(f"[engram-debug] engram_out: finite={torch.isfinite(engram_out).all().item()} min={engram_out.float().min().item():.6g} max={engram_out.float().max().item():.6g}", flush=True) + engram_out = torch.nan_to_num(engram_out, nan=0.0, posinf=0.0, neginf=0.0).clamp_(-100.0, 100.0) + + # alpha_mean: mean gate activity — used by LatentMemory for suppression + alpha_mean = alpha_t.squeeze(-1) # [B, T] + + return engram_out, alpha_mean diff --git a/fla/__init__.py b/fla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d545d1c14c78c2c897030200505e22f4eb27d640 --- /dev/null +++ b/fla/__init__.py @@ -0,0 +1,8 @@ +# Lightweight package initializer for the vendored FLA subset used by Quasar. +# +# The upstream file eagerly imports every layer and model, which is slow and can +# hang on fresh training containers while optional kernels are being resolved. +# Import concrete modules directly, e.g. `from fla.layers.quasar import ...`. + +__version__ = "0.1.0" +__all__ = [] diff --git a/fla/distributed_compat.py b/fla/distributed_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..b31933781302283d424cfa5b28e84fb3990068c0 --- /dev/null +++ b/fla/distributed_compat.py @@ -0,0 +1,57 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +""" +Centralized compatibility module for torch.distributed imports. +All distributed-related imports should go through here to handle environments +where distributed tensor APIs are not available. +""" + +import torch + +# DeviceMesh +try: + from torch.distributed import DeviceMesh +except ImportError: + try: + from torch.distributed.device_mesh import DeviceMesh + except ImportError: + DeviceMesh = None + +# DTensor +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + +# Replicate, Shard, distribute_module, Placement +try: + from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module +except (ImportError, AttributeError): + Placement = Replicate = Shard = distribute_module = None + +# ParallelStyle +try: + from torch.distributed.tensor.parallel import ParallelStyle +except (ImportError, AttributeError): + ParallelStyle = None + +# Convenience flag +HAS_DISTRIBUTED = all([ + DeviceMesh is not None, + DTensor is not None, + Placement is not None, + Replicate is not None, + Shard is not None, + distribute_module is not None, + ParallelStyle is not None, +]) + +__all__ = [ + 'DeviceMesh', + 'DTensor', + 'Placement', + 'Replicate', + 'Shard', + 'distribute_module', + 'ParallelStyle', + 'HAS_DISTRIBUTED', +] diff --git a/fla/layers/__init__.py b/fla/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d68fde5a2384108df341419ac1a0bff6d455d640 --- /dev/null +++ b/fla/layers/__init__.py @@ -0,0 +1,3 @@ +# Keep layer package imports lazy. Import specific layer modules directly. + +__all__ = [] diff --git a/fla/layers/abc.py b/fla/layers/abc.py new file mode 100644 index 0000000000000000000000000000000000000000..99469a7af40c2213970ff1f7fbc523f8b0df28e2 --- /dev/null +++ b/fla/layers/abc.py @@ -0,0 +1,231 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, RotaryEmbedding, ShortConvolution +from fla.modules.activations import swiglu, swish +from fla.ops.abc.chunk import chunk_abc + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class ABCAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 1024, + expand_k: float = 0.5, + expand_v: float = 1.0, + num_heads: int = 4, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_slots: int | None = None, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + gate_low_rank_dim: int = 16, + gate_logit_normalizer: int = 16, + use_rope: bool = True, + use_input_gate: bool = False, + use_output_gate: bool = True, + use_norm: bool = True, + clamp_min: float | None = -32, + clamp_max: float | None = 32, + layer_idx: int | None = None, + **kwargs, + ) -> ABCAttention: + super().__init__() + + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.key_dim = int(self.hidden_size * self.expand_k) + self.value_dim = int(self.hidden_size * self.expand_v) + self.head_k_dim = self.key_dim // self.num_heads + self.head_v_dim = self.value_dim // self.num_heads + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.gate_low_rank_dim = gate_low_rank_dim + self.gate_logit_normalizer = gate_logit_normalizer + + self.use_rope = use_rope + self.use_input_gate = use_input_gate + self.use_output_gate = use_output_gate + self.use_norm = use_norm + + if num_slots is None: + num_slots = self.head_k_dim + self.num_slots = num_slots + + self.norm_eps = norm_eps + + self.clamp_min = clamp_min + self.clamp_max = clamp_max + self.layer_idx = layer_idx + + if layer_idx is None: + warnings.warn( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class.", + ) + + self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False) + + if use_output_gate: + self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False) + self.s_proj = nn.Linear(self.hidden_size, self.num_heads * self.num_slots, bias=False) + self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + if self.use_norm: + if self.use_output_gate: + self.g_norm = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + else: + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + dtype=torch.float32, + ) + + if self.use_rope: + self.rotary = RotaryEmbedding(self.head_k_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if cu_seqlens is not None: + raise NotImplementedError("Training with cu_seqlens is not supported yet for ABCAttention") + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + if self.use_input_gate: + q, k, v = map(lambda x: swish(x), (q, k, v)) + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + if self.use_rope: + seqlen_offset = 0 + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset) + + s = rearrange(self.s_proj(hidden_states), '... (h m) -> ... h m', m=self.num_slots) + s = s.clamp_(self.clamp_min, self.clamp_max) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + o, recurrent_state = chunk_abc( + q=q, + k=k, + v=v, + s=s, + initial_state=recurrent_state, + output_final_state=use_cache, + ) + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q.shape[1], + ) + + if self.use_norm and not self.use_output_gate: + o = self.g_norm(o) + elif self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.g_norm(o, g) if self.use_norm else swiglu(g, o) + o = rearrange(o, '... h d -> ... (h d)') + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, seq_len: int = 2048): + return 2 * self.num_slots * self.hidden_size diff --git a/fla/layers/attn.py b/fla/layers/attn.py new file mode 100644 index 0000000000000000000000000000000000000000..658f023a43c483042ed87ada956c3b72e2977733 --- /dev/null +++ b/fla/layers/attn.py @@ -0,0 +1,176 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.utils import logging + +from fla.layers.utils import pad_input, unpad_input +from fla.modules import RMSNorm, RotaryEmbedding +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +class Attention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: int | None = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: int | None = None, + rope_theta: float | None = 10000., + max_position_embeddings: int | None = None, + use_nope: bool = False, + layer_idx: int = None, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.use_nope = use_nope + self.layer_idx = layer_idx + + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32) + self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + # equivalent to cu_seqlens in `flash_attn` + cu_seqlens = kwargs.get('cu_seqlens') + + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + if not self.use_nope: + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size), + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + # Contains at least one padding token in the sequence + if attention_mask is not None: + if q.shape[1] == 1 and self.window_size is not None: + attention_mask = attention_mask[:, -self.window_size:] + q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len) + cu_seqlens_q, cu_seqlens_k = cu_seqlens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = pad_input(o, indices_q, batch_size, q_len) + elif cu_seqlens is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ).unsqueeze(0) + else: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/based.py b/fla/layers/based.py new file mode 100644 index 0000000000000000000000000000000000000000..fe1adfebdc0c2ca28c860ff7fff942ae8238b082 --- /dev/null +++ b/fla/layers/based.py @@ -0,0 +1,93 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +Linear attention in Based. +https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py +""" + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.modules.feature_map import TaylorFeatureMap +from fla.ops.based import parallel_based +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn + + +class BasedLinearAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + feature_dim: int = 16, + num_key_value_heads: int = 12, + num_heads: int = 12, + feature_name: str = "taylor_exp", + eps: float = 1e-12, + causal: bool = True, + mode: str = "parallel", + ): + super().__init__() + + self.hidden_size = hidden_size + self.mode = mode + self.feature_name = feature_name + self.feature_dim = feature_dim + self.num_key_value_heads = num_key_value_heads + self.num_heads = num_heads + self.head_dim = self.hidden_size // self.num_key_value_heads + assert self.hidden_size % self.head_dim == 0 + self.causal = causal + + self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.dropout = nn.Identity() + self.feature_map = TaylorFeatureMap(feature_dim) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor, **kwargs): + mode = self.mode + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v]) + if mode == "fused_chunk": + q, k = self.feature_map(q), self.feature_map(k) + o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1) + elif mode == 'chunk': + q, k = self.feature_map(q), self.feature_map(k) + o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1) + elif mode == 'parallel': + assert q.shape[-1] <= 128 + o = parallel_based(q, k, v, scale=1, use_norm=True) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + o = self.dropout(o) + return o + + def forward_reference(self, hidden_states: torch.Tensor, **kwargs): + """ + x (torch.Tensor): tensor of shape (b, d, t) + y (torch.Tensor): tensor of shape (b, d, t) + """ + # hidden_states = hidden_states.transpose(1, 2) + b, t, _ = hidden_states.size() + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + + q = q.view(b, t, self.num_heads, self.feature_dim).transpose(1, 2) + k = k.view(b, t, self.num_key_value_heads, self.feature_dim).transpose(1, 2) + v = v.view(b, t, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + # Linear attention + q, k = self.feature_map(q), self.feature_map(k) + q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1) + + # Compute attention + if self.causal: + y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps)) + else: + y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps)) + y = rearrange(y, 'b h t d -> b t (h d)') + y = self.o_proj(y.to(hidden_states.dtype)) + y = self.dropout(y) + return y.to(hidden_states.dtype) diff --git a/fla/layers/bitattn.py b/fla/layers/bitattn.py new file mode 100644 index 0000000000000000000000000000000000000000..40d0aeff0989b9f65216b7bd3350de4e6991c06a --- /dev/null +++ b/fla/layers/bitattn.py @@ -0,0 +1,162 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.utils import logging + +from fla.layers.utils import pad_input, unpad_input +from fla.modules import RotaryEmbedding +from fla.modules.fused_bitlinear import FusedBitLinear +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +class BitAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: int | None = None, + window_size: int | None = None, + rope_theta: float | None = 10000., + max_position_embeddings: int | None = None, + norm_eps: float = 1e-5, + layer_idx: int = None, + ): + super().__init__() + + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.hidden_size = hidden_size + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False) + self.k_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False) + self.v_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False) + self.o_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + + # equivalent to cu_seqlens in `flash_attn` + cu_seqlens = kwargs.get('cu_seqlens') + + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size), + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + # Contains at least one padding token in the sequence + if attention_mask is not None: + q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len) + cu_seqlens_q, cu_seqlens_k = cu_seqlens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = pad_input(o, indices_q, batch_size, q_len) + elif cu_seqlens is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ).unsqueeze(0) + else: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/comba.py b/fla/layers/comba.py new file mode 100644 index 0000000000000000000000000000000000000000..a253ff2eae47c8913f53c2d1722281a7155a00e1 --- /dev/null +++ b/fla/layers/comba.py @@ -0,0 +1,328 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.comba import chunk_comba, fused_recurrent_comba + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class Comba(nn.Module): + """ + The layer implementaion for [Comba: Improving Bilinear RNNs with Closed-loop Control](https://arxiv.org/abs/2506.02475). + + Similar to Mamba2 and Gated-DeltaNet, each layer contains around 6*hidden_size*hidden_size parameters. + + Parameter alloation when use_output_gate=True: + - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each + - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each + - Others are ignorably small. + - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size + NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim. + + Parameter allocation when use_output_gate=False: + - 1 * hidden_size * hidden_size for the q_proj and k_proj each + - 2 * hidden_size * hidden_size for the v_proj and o_proj each + - Others are ignorably small. + - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 4. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_beta (bool, Optional): + Whether to use beta. Default: `True`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_output_correction (bool, Optional): + Whether to use . Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_short_conv: bool = True, + use_output_gate: bool = True, + use_output_correction: bool = True, + use_inner_decay: bool = True, + correction_factor: float = 1., + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> Comba: + super().__init__() + + self.mode = mode + + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_short_conv = use_short_conv + self.use_output_gate = use_output_gate + self.use_output_correction = use_output_correction + self.use_inner_decay = use_inner_decay + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + + if use_inner_decay: + self.decay = nn.Parameter(torch.ones(self.num_heads)) + + if use_output_correction: + warnings.warn( + "The correction_factor is set to 1 by default similar to Mamba2. " + "However, we find that sometimes correction_factor = 0.02 works better for small-scale models. " + "In practice, we recommend trying both settings. ", + ) + self.D = nn.Parameter(torch.ones(self.num_heads) * correction_factor) + self.D._no_weight_decay = True + + A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + warnings.warn( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, activation='sigmoid', eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + + if self.use_inner_decay: + p = k * self.decay[None, None, :, None].sigmoid() + else: + p = k + + if self.use_output_correction: + q = q - self.D[None, None, :, None] * p + + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + if self.num_v_heads > self.num_heads: + q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k)) + + beta = self.b_proj(hidden_states).sigmoid() + g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_comba( + q=q, + k=k, + v=v, + p=p, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_comba( + q=q, + k=k, + v=v, + p=p, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/delta_net.py b/fla/layers/delta_net.py new file mode 100644 index 0000000000000000000000000000000000000000..26928e16e12da1280e185c4ec19f04e1ef3385fd --- /dev/null +++ b/fla/layers/delta_net.py @@ -0,0 +1,287 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +def elu_p1(x): + return (F.elu(x, 1., False) + 1.).to(x) + + +def sum_norm(x): + return (x / x.sum(-1, keepdim=True)).to(x) + + +class DeltaNet(nn.Module): + r""" + The layer implementaion for [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://arxiv.org/abs/2406.06484). # noqa: + DeltaNet was originally proposed in [Linear Transformers Are Secretly Fast Weight Programmers](https://arxiv.org/abs/2102.11174). # noqa + + Args: + mode (str, Optional): + Which DeltaNet kernel to use. + Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + use_beta (bool, Optional): + Whether to use beta. Default: `True`. + use_gate (bool, Optional): + Whether to use output gate. Default: `False`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + qk_activation (str, Optional): + The activation function for the query and key. Default: `silu`. + qk_norm (str, Optional): + The normalization method for the query and key. Default: `l2`. + """ + + def __init__( + self, + mode: str = 'chunk', + d_model: int = None, + hidden_size: int = 1024, + expand_k: float = 1.0, + expand_v: float = 1.0, + num_heads: int = 4, + use_beta: bool = True, + use_gate: bool = False, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + allow_neg_eigval: bool = False, + layer_idx: int = None, + qk_activation: str = 'silu', + qk_norm: str = 'l2', + norm_eps: float = 1e-5, + **kwargs, + ) -> DeltaNet: + super().__init__() + + self.mode = mode + self.qk_activation = qk_activation + self.qk_norm = qk_norm + + assert self.qk_activation in ['silu', 'relu', 'elu', 'identity'] + assert self.qk_norm in ['l2', 'sum'] + + if d_model is not None: + hidden_size = d_model + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.allow_neg_eigval = allow_neg_eigval + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + self.layer_idx = layer_idx + + if mode == 'fused_chunk': + raise NotImplementedError("fused_chunk_delta_rule is now deprecated. Please use `chunk_delta_rule` instead.") + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + self.use_beta = use_beta + if self.use_beta: + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu' if qk_activation == 'silu' else None, + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu' if qk_activation == 'silu' else None, + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + warnings.warn( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + if self.qk_activation == 'silu': + q, k = F.silu(q), F.silu(k) + v = F.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + if self.qk_activation != 'silu': + if self.qk_activation == 'relu': + q, k = q.relu(), k.relu() + elif self.qk_activation == 'elu': + q, k = elu_p1(q), elu_p1(k) + elif self.qk_activation != 'identity': + raise NotImplementedError + + if self.qk_norm == 'sum': + q = sum_norm(q).to(q) + k = sum_norm(k).to(k) + + if self.use_beta: + beta = self.b_proj(hidden_states).sigmoid() + else: + beta = torch.ones_like(q[..., 0]) + + if self.allow_neg_eigval: + beta = beta * 2. + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=(self.qk_norm == 'l2'), + ) + elif mode == 'chunk': + o, recurrent_state = chunk_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=(self.qk_norm == 'l2'), + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/deltaformer.py b/fla/layers/deltaformer.py new file mode 100644 index 0000000000000000000000000000000000000000..e7e62af781cee40b372985160b5b847c2f96c230 --- /dev/null +++ b/fla/layers/deltaformer.py @@ -0,0 +1,152 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.utils import logging + +from fla.modules import RMSNorm, RotaryEmbedding +from fla.ops.deltaformer import deltaformer_attn +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class DeltaFormerAttention(nn.Module): + + r""" + The layer implementation for DeltaFormer, + [Understanding Transformer from the Perspective of Associative Memory] + (https://arxiv.org/pdf/2505.19488). + + Notes + - DeltaFormer attention is implemented with Triton kernels in `fla.ops.deltaformer` and is tuned + for typical head dimensions (e.g., 64/128). It currently supports fixed-length inputs. + - For variable-length inputs (padding masks), the deltaformer computation falls back to using the + fixed-length path, while the second stage (softmax attention over U) uses FlashAttention's + varlen path when an attention mask is provided. + - K/V grouping (GQA) is supported natively by FlashAttention via `num_kv_heads`. + - Uses K-K similarity in deltaformer computation instead of Q-K similarity for better performance. + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + num_heads (int, Optional): + The number of attention heads. Default: 32. + num_kv_heads (int, Optional): + The number of key/value heads for grouped-query attention. If None, equals `num_heads`. + Default: None. + qkv_bias (bool, Optional): + Whether to use bias for Q/K/V projections. Default: False. + qk_norm (bool, Optional): + Whether to apply per-head RMSNorm to Q and K before attention. Default: False. + rope_theta (float, Optional): + The base frequency for rotary position embedding. Default: 10000. + max_position_embeddings (int, Optional): + The maximum position embeddings. Default: None. + layer_idx (int, Optional): + The index of the layer (used for cache compatibility). Default: None. + """ + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: int | None = None, + qkv_bias: bool = False, + qk_norm: bool = False, + rope_theta: float = 10000., + max_position_embeddings: int | None = None, + layer_idx: int | None = None, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32) + self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + attentions = None + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + beta = self.b_proj(hidden_states) + + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + cu_seqlens_kw = kwargs.get('cu_seqlens') + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q_len + seqlen_offset + + if attention_mask is not None: + seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1] + max_seqlen = q_len + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens_kw) + + o = deltaformer_attn( + q=q, + k=k, + v=v, + beta=beta, + attention_mask=attention_mask, + cu_seqlens=cu_seqlens_kw, + ) + + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/forgetting_attn.py b/fla/layers/forgetting_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..5dcf9cd62c997cc7b16ffcc4a4ff1fbf7c029392 --- /dev/null +++ b/fla/layers/forgetting_attn.py @@ -0,0 +1,133 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange +from transformers.utils import logging + +from fla.layers.utils import pad_input, unpad_input +from fla.modules import GroupNorm +from fla.ops.attn.decoding import attn_decoding_one_step +from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class ForgettingAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: int | None = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: int | None = None, + use_output_gate: bool = False, + layer_idx: int = None, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + + self.window_size = window_size + self.use_output_gate = use_output_gate + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.f_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True) + + if use_output_gate: + self.g_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = GroupNorm( + num_groups=self.num_heads, + hidden_size=self.hidden_size, + is_rms_norm=True, + ) + self.k_norm = GroupNorm( + num_groups=self.num_kv_heads, + hidden_size=self.kv_dim, + is_rms_norm=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + f = F.logsigmoid(self.f_proj(hidden_states).float()) + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + cu_seqlens = kwargs.get('cu_seqlens') + if past_key_values is not None: + assert cu_seqlens is None, "cu_seqlens should not be provided when past_key_values is not None" + state = past_key_values.update( + attn_state=(k, v, f), + layer_idx=self.layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size), + ) + k, v, f = state['attn_state'] + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if attention_mask is not None: + q, (k, v, f), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v, f), attention_mask, q_len, keepdim=True) + _, cu_seqlens_k = cu_seqlens + cu_seqlens = cu_seqlens_k + max_seqlen_q, max_seqlen_k = max_seq_lens + if max_seqlen_q != max_seqlen_k: + assert max_seqlen_q == 1, "only support q_len == 1 for decoding" + o = attn_decoding_one_step(q, k, v, f, cu_seqlens=cu_seqlens) + else: + o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens) + else: + o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices_q, batch_size, q_len) + o = rearrange(o, '... h d -> ... (h d)') + if self.use_output_gate: + o = self.g_proj(hidden_states).sigmoid() * o + o = self.o_proj(o) + return o, None, past_key_values diff --git a/fla/layers/gated_deltanet.py b/fla/layers/gated_deltanet.py new file mode 100644 index 0000000000000000000000000000000000000000..967724bb72830b9417283dfe7f2e946c20b4fa2f --- /dev/null +++ b/fla/layers/gated_deltanet.py @@ -0,0 +1,316 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +@torch.compile +def elu_p1(x): + return (F.elu(x, 1., False) + 1.).to(x) + + +@torch.compile +def sum_norm(x): + return (x / x.sum(-1, keepdim=True)).to(x) + + +class GatedDeltaNet(nn.Module): + """ + The layer implementaion for [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://arxiv.org/abs/2412.06464). # noqa + + Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters. + + Parameter alloation when use_gate=True: + - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each + - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each + - Others are ignorably small. + - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size + NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim. + + Parameter allocation when use_gate=False: + - 1 * hidden_size * hidden_size for the q_proj and k_proj each + - 2 * hidden_size * hidden_size for the v_proj and o_proj each + - Others are ignorably small. + - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 4. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_beta (bool, Optional): + Whether to use beta. Default: `True`. + use_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_gate: bool = True, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> GatedDeltaNet: + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + + A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + warnings.warn( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + if self.num_v_heads > self.num_heads: + q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k)) + + beta = self.b_proj(hidden_states).sigmoid() + if self.allow_neg_eigval: + beta = beta * 2. + + g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/gated_deltaproduct.py b/fla/layers/gated_deltaproduct.py new file mode 100644 index 0000000000000000000000000000000000000000..01195ebda672db789baefd581921defd0de5921a --- /dev/null +++ b/fla/layers/gated_deltaproduct.py @@ -0,0 +1,287 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.gated_delta_product import chunk_gated_delta_product +from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedDeltaProduct(nn.Module): + """ + Generalized version of GatedDoubleDeltaNet that supports arbitrary number of householder transformations. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_output_gate: bool = True, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + use_forget_gate: bool = True, + allow_neg_eigval: bool = True, + num_householder: int = 2, + **kwargs, + ) -> GatedDeltaProduct: + super().__init__() + + self.mode = mode + + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_forget_gate = use_forget_gate + self.allow_neg_eigval = allow_neg_eigval + self.num_householder = num_householder + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim * num_householder, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim * num_householder, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads * num_householder, bias=False) + + if self.use_forget_gate: + self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim * num_householder, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim * num_householder, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + warnings.warn( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5) + if module.bias is not None: + nn.init.zeros_(module.bias) + module._is_hf_initialized = True + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + k = rearrange(k, '... t (n h d) -> ... (t n) h d', n=self.num_householder, d=self.head_k_dim) + v = rearrange(v, '... t (n h d) -> ... (t n) h d', n=self.num_householder, d=self.head_v_dim) + + if self.num_v_heads > self.num_heads: + q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k)) + + beta = self.b_proj(hidden_states).sigmoid() + if self.allow_neg_eigval: + beta = beta * 2. + + beta = rearrange(beta, '... t (n h) -> ... (t n) h', n=self.num_householder) + if self.use_forget_gate: + g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) + else: + g = None + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_gated_delta_product( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + num_householder=self.num_householder, + use_qk_l2norm_in_kernel=True, + ) + + elif mode == 'fused_recurrent': + if self.use_forget_gate: + g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2]) + g_new[:, :, 0] = g + g = rearrange(g_new, '... t n h -> ... (t n) h') + + q_new = q.new_zeros(q.shape[0], q.shape[1], self.num_householder, q.shape[2], q.shape[3]) + q_new[:, :, -1] = q + q = rearrange(q_new, '... t n h d-> ... (t n) h d') + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens * self.num_householder if cu_seqlens is not None else None, + use_qk_l2norm_in_kernel=True, + ) + o = rearrange(o, '... (t n) h d -> ... t n h d', n=self.num_householder)[..., -1, :, :].contiguous() + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/gla.py b/fla/layers/gla.py new file mode 100644 index 0000000000000000000000000000000000000000..8dac3b21e519467ddc1da130f3782927c62cb965 --- /dev/null +++ b/fla/layers/gla.py @@ -0,0 +1,304 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.activations import ACT2FN +from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedLinearAttention(nn.Module): + r""" + The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635). # noqa + + Args: + mode (str, Optional): + Which GLA kernel to use. + Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 0.5. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + gate_logit_normalizer (int, Optional): + The normalizer for the gate logits, appied after `logsigmoid`. Default: 16. + gate_low_rank_dim (int, Optional): + The low rank dim for the gate projection. Default: 16. + clamp_min (float, Optional): + The minimum value for the gate logits. Default: None. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 0.5, + expand_v: float = 1.0, + num_heads: int = 4, + num_kv_heads: int | None = None, + feature_map: str | None = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + use_output_gate: bool = True, + gate_fn: str = 'swish', + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 16, + gate_low_rank_dim: int = 16, + clamp_min: float | None = None, + fuse_norm: bool = True, + layer_idx: int = None, + ) -> GatedLinearAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.use_output_gate = use_output_gate + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.clamp_min = clamp_min + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + if self.use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + self.gk_proj = nn.Sequential(nn.Linear(hidden_size, gate_low_rank_dim, bias=False), + nn.Linear(gate_low_rank_dim, self.key_dim_per_group, bias=True)) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + if gate_fn == 'swish' and fuse_norm and use_output_gate: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + dtype=torch.float32 + ) + self.gate_fn = ACT2FN[gate_fn] + + self.gate_logit_normalizer = gate_logit_normalizer + + def reset_parameters(self) -> None: + for module in self.children(): + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + gk = self.gk_proj(hidden_states) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + if self.num_kv_groups > 1: + k, gk = (repeat(x, '... (h d) -> ... (h g) d', g=self.num_kv_groups, d=self.head_k_dim) for x in (k, gk)) + v = repeat(v, '... (h d) -> ... (h g) d', g=self.num_kv_groups, d=self.head_v_dim) + else: + k, gk = (rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim) for x in (k, gk)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + gk = F.logsigmoid(gk) / self.gate_logit_normalizer + if self.clamp_min is not None: + gk = torch.clamp_min(gk, self.clamp_min) + + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_output_gate: + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, '... (h d) -> ... h d', d=self.head_v_dim) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, '... h d -> ... (h d)') + else: + o = rearrange(self.g_norm(o), '... h d -> ... (h d)') + o = o * self.gate_fn(g) + else: + o = rearrange(self.g_norm(o), '... h d -> ... (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/gsa.py b/fla/layers/gsa.py new file mode 100644 index 0000000000000000000000000000000000000000..73a95e3029b2db6822ad3d9359d77480f173c73a --- /dev/null +++ b/fla/layers/gsa.py @@ -0,0 +1,237 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import RMSNorm, ShortConvolution +from fla.modules.feature_map import ReLUFeatureMap, SwishFeatureMap, T2RFeatureMap +from fla.modules.layernorm import rms_norm_linear +from fla.ops.gsa import chunk_gsa, fused_recurrent_gsa + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedSlotAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1., + expand_v: float = 1., + num_heads: int = 4, + num_kv_heads: int | None = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_slots: int | None = None, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 8, + feature_map: str = 'swish', + use_output_gate: bool = False, + use_norm: bool = True, + layer_idx: int | None = None, + scale: float | None = 1., + **kwargs, + ) -> GatedSlotAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.head_k_dim = self.key_dim // self.num_heads + self.head_v_dim = self.value_dim // self.num_heads + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.gate_logit_normalizer = gate_logit_normalizer + + self.use_output_gate = use_output_gate + self.use_norm = use_norm + self.scale = scale + + if num_slots is None: + num_slots = self.head_k_dim + self.num_slots = num_slots + + self.layer_idx = layer_idx + + if layer_idx is None: + warnings.warn( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class.", + ) + + self.register_module('feature_map', None) + if feature_map == 'swish': + self.feature_map = SwishFeatureMap() + elif feature_map == 'relu': + self.feature_map = ReLUFeatureMap() + elif feature_map == 't2r': + self.feature_map = T2RFeatureMap(self.head_k_dim, self.head_k_dim) + else: + raise NotImplementedError(f"Feature map `{feature_map}` is not supported now.") + + self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.value_dim_per_group, bias=False) + self.f_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.num_slots, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + self.g_norm = RMSNorm(self.hidden_size, elementwise_affine, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + f = self.f_proj(hidden_states) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + f = rearrange(f, '... (h m) -> ... h m', m=self.num_slots) + + if self.feature_map is not None: + q, k = map(lambda x: self.feature_map(x), (q, k)) + v = F.silu(v) + + f = F.logsigmoid(f) / self.gate_logit_normalizer + s = (1 - f.exp()).to(f.dtype) + + if self.num_kv_groups > 1: + k, v, f, s = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_kv_groups), (k, v, f, s)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + scale=self.scale, + cu_seqlens=cu_seqlens, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gsa( + q=q, + k=k, + v=v, + s=s, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + scale=self.scale, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + o = rearrange(o, '... h d -> ... (h d)') + o = rms_norm_linear(F.silu(o), self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/hgrn.py b/fla/layers/hgrn.py new file mode 100644 index 0000000000000000000000000000000000000000..92cfd3adeb6075e0b333878959dad5d00b99bb10 --- /dev/null +++ b/fla/layers/hgrn.py @@ -0,0 +1,174 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# "Hierarchically Gated Recurrent Neural Network for Sequence Modeling" [https://arxiv.org/abs/2311.04823] + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.modules.activations import swiglu +from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class HGRNAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_ratio: int | None = 1, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + ) -> HGRNAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_ratio = expand_ratio + self.input_dim = int(hidden_size * expand_ratio) + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + self.f_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + self.g_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.f_conv1d = ShortConvolution( + hidden_size=self.input_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.i_conv1d = ShortConvolution( + hidden_size=self.input_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + + self.g_norm = FusedRMSNormGated( + hidden_size=self.input_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + lower_bound: torch.Tensor | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if not self.training and hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if self.use_short_conv: + conv_state_i, conv_state_f = None, None + if last_state is not None: + conv_state_i, conv_state_f = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + i, conv_state_i = self.i_conv1d( + x=self.i_proj(hidden_states), + mask=conv_mask, + cache=conv_state_i, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + f, conv_state_f = self.f_conv1d( + x=self.f_proj(hidden_states), + mask=conv_mask, + cache=conv_state_f, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + i = self.i_proj(hidden_states) + f = self.f_proj(hidden_states) + + f = F.logsigmoid(f) + # the lower bound for the first layer is zero + if lower_bound is not None and self.layer_idx > 0: + f = torch.logaddexp(lower_bound.log(), torch.log1p(-lower_bound) + f).to(f) + i = swiglu(i, 1 - f.exp()) + + # dealing with left-padding + if attention_mask is not None: + i = i.mul(attention_mask[:, -i.shape[-2]:, None]) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + if cu_seqlens is not None: + raise NotImplementedError("Chunk mode does not support variable-length sequences.") + o, recurrent_state = chunk_hgrn( + x=i, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_hgrn( + x=i, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_i, conv_state_f) if self.use_short_conv else None, + offset=i.shape[1], + ) + + o = self.g_norm(o, self.g_proj(hidden_states)) + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.hidden_size + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/hgrn2.py b/fla/layers/hgrn2.py new file mode 100644 index 0000000000000000000000000000000000000000..8596fe4cef6f7da74e541004ca76cb6588ed15c5 --- /dev/null +++ b/fla/layers/hgrn2.py @@ -0,0 +1,210 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# "HGRN2: Gated Linear RNNs with State Expansion"[https://arxiv.org/abs/2404.07904] + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import RMSNorm, ShortConvolution +from fla.modules.activations import swish +from fla.modules.layernorm import rms_norm_linear +from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class HGRN2Attention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + num_heads: int | None = None, + expand_ratio: int | None = 128, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + ) -> HGRN2Attention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + + if expand_ratio is not None: + num_heads = hidden_size // expand_ratio + elif expand_ratio is None and num_heads is not None: + expand_ratio = hidden_size // num_heads + elif expand_ratio is None and num_heads is None: + raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.") + self.num_heads = num_heads + self.expand_ratio = expand_ratio + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.forget_dim = int(self.num_heads * self.expand_ratio) + self.input_dim = hidden_size + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`." + assert self.forget_dim % num_heads == 0, f"forget dim must be divisible by num_heads of {num_heads}" + assert self.input_dim % num_heads == 0, f"input dim must be divisible by num_heads of {num_heads}" + + self.head_f_dim = self.expand_ratio + self.head_i_dim = self.hidden_size // num_heads + + self.q_proj = nn.Linear(hidden_size, self.forget_dim, bias=False) + self.f_proj = nn.Linear(hidden_size, self.forget_dim, bias=False) + self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.forget_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.f_conv1d = ShortConvolution( + hidden_size=self.forget_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.i_conv1d = ShortConvolution( + hidden_size=self.input_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + + self.g_norm = RMSNorm(hidden_size=self.hidden_size, elementwise_affine=elementwise_affine, + eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + lower_bound: torch.Tensor | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_f, conv_state_i = None, None, None + if last_state is not None: + conv_state_q, conv_state_f, conv_state_i = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + f, conv_state_f = self.f_conv1d( + x=self.f_proj(hidden_states), + cache=conv_state_f, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + i, conv_state_i = self.i_conv1d( + x=self.i_proj(hidden_states), + cache=conv_state_i, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + f = self.f_proj(hidden_states) + i = self.i_proj(hidden_states) + + q = swish(q) + + g = F.logsigmoid(f) + # the lower bound for the first layer is zero + if lower_bound is not None and self.layer_idx > 0: + g = torch.logaddexp(lower_bound.log(), torch.log1p(-lower_bound) + g) + k = 1 - g.exp() + + q, k, g = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k.to(i), g)) + i = rearrange(i, '... (h d) -> ... h d', d=self.head_i_dim) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=i, + gk=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_gla( + q=q, + k=k, + v=i, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=i, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_f, conv_state_i) if self.use_short_conv else None, + offset=q_len, + ) + + o = rearrange(o, '... h d -> ... (h d)') + o = rms_norm_linear(o, self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/kda.py b/fla/layers/kda.py new file mode 100644 index 0000000000000000000000000000000000000000..85c06221b80bff9c24b052050e9dc0c66f07612b --- /dev/null +++ b/fla/layers/kda.py @@ -0,0 +1,277 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.ops.kda import chunk_kda, fused_recurrent_kda +from fla.ops.kda.gate import fused_kda_gate + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class KimiDeltaAttention(nn.Module): + """ + Kimi Delta Attention (KDA) layer implementation. + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dimension. Default: 1.0. + head_dim (int, Optional): + The dimension of each head. Default: 128. + num_heads (int, Optional): + The number of heads. Default: 16. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA (Grouped Value Attention) is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which Kimi Delta Attention kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: + [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 1, + head_dim: int = 128, + num_heads: int = 16, + num_v_heads: int = None, + mode: str = "chunk", + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> KimiDeltaAttention: + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + + self.f_proj = nn.Sequential( + nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=False), + ) + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + + self.A_log = nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16))) + self.A_log._no_weight_decay = True + dt = torch.exp( + torch.rand(self.key_dim, dtype=torch.float32) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) + ).clamp(min=1e-4) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + self.dt_bias._no_weight_decay = True + + self.g_proj = nn.Sequential( + nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.value_dim, bias=True), + ) + self.o_norm = FusedRMSNormGated(self.head_v_dim, activation="sigmoid", eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = "fused_recurrent" if (q_len <= 64 and not self.training) else self.mode + if self.training: + assert mode == "chunk", "Only chunk mode is supported in training." + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get("cu_seqlens") + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + g = self.f_proj(hidden_states) + beta = self.b_proj(hidden_states).sigmoid() + + q, k, g = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) for x in (q, k, g)) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_v_dim) + + # for multi-value attention, we repeat the inputs for simplicity. + if self.num_v_heads > self.num_heads: + q, k, g = (repeat(x, "... h d -> ... (h g) d", g=self.num_v_heads // self.num_heads) for x in (q, k, g)) + beta = repeat(beta, "... h -> ... (h g)", g=self.num_v_heads // self.num_heads) + + if self.allow_neg_eigval: + beta = beta * 2.0 + + recurrent_state = last_state["recurrent_state"] if last_state is not None else None + if mode == "chunk": + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + elif mode == "fused_recurrent": + g = fused_kda_gate(g=g, A_log=self.A_log, dt_bias=self.dt_bias) + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + o = self.o_norm(o, rearrange(self.g_proj(hidden_states), "... (h d) -> ... h d", d=self.head_v_dim)) + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/lightnet.py b/fla/layers/lightnet.py new file mode 100644 index 0000000000000000000000000000000000000000..0a559e185f445b9adeb23c64a894c902e44ee009 --- /dev/null +++ b/fla/layers/lightnet.py @@ -0,0 +1,238 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# ["You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet"](https://arxiv.org/abs/2405.21022) + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.modules.fused_norm_gate import rms_norm_swish_gate_linear +from fla.ops.gla import chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class LightNetAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + num_heads: int | None = None, + expand_ratio: int | None = 128, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + gate_low_rank_dim: int = 128, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + ) -> LightNetAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + + if expand_ratio is None and num_heads is not None: + expand_ratio = hidden_size // num_heads + elif expand_ratio is not None and num_heads is None: + num_heads = hidden_size // expand_ratio + elif expand_ratio is None and num_heads is None: + raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.") + self.num_heads = num_heads + self.expand_ratio = expand_ratio + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.key_dim = int(self.num_heads * self.expand_ratio) + self.value_dim = hidden_size + self.gate_low_rank_dim = gate_low_rank_dim + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_chunk'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_f_dim = self.expand_ratio + self.head_i_dim = self.hidden_size // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + + self.g_proj = nn.Sequential( + nn.Linear(hidden_size, gate_low_rank_dim, bias=False), + nn.Linear(gate_low_rank_dim, hidden_size, bias=False), + ) + self.g_norm = FusedRMSNormGated( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # dealing with left-padding + if attention_mask is not None: + v = v.mul(attention_mask[:, -v.shape[-2]:, None]) + + q = F.silu(q) + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_i_dim) + # TODO: this 2 steps took huge amount of time, which should be optimized + last_z = last_state['ffn_state'] if last_state is not None and last_state.get('ffn_state') is not None else None + if last_z is not None: + # Decode path: continue logcumsumexp from cached state + z = torch.logaddexp(last_z, k.float()) + k, g = torch.exp(k - z).to(k.dtype), (last_z - z).to(k.dtype) + else: + # Prefill path: mask padding positions to -inf so they don't affect logcumsumexp + if cu_seqlens is not None: + raise NotImplementedError("LightNet does not support variable-length sequences for now.") + k_float = k.float() + if attention_mask is not None: + pad_mask = attention_mask[:, -k.shape[1]:, None, None] # (B, T, 1, 1) + k_for_z = k_float.masked_fill(pad_mask == 0, float('-inf')) + else: + k_for_z = k_float + z = k_for_z.logcumsumexp(1) + k_new = torch.exp(k_float - z) + g_new = torch.cat((z[:, :1], z[:, :-1]), 1) - z + # NaN/inf arise at fully-masked positions (-inf - (-inf)), zero them out + k = torch.nan_to_num(k_new, nan=0.0, posinf=0.0).to(k.dtype) + g = torch.nan_to_num(g_new, nan=0.0, posinf=0.0, neginf=0.0).to(k.dtype) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + ffn_state=z[:, -1:], + offset=q.shape[1], + ) + + o = rms_norm_swish_gate_linear( + rearrange(o, 'b t h d -> b t (h d)'), + self.g_proj(hidden_states), + self.g_norm.weight, + self.g_norm.bias, + self.o_proj.weight, + self.o_proj.bias, + ) + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_i_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/linear_attn.py b/fla/layers/linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..dc3c8f33527b9eaa1ac535844f2d54713527fa6d --- /dev/null +++ b/fla/layers/linear_attn.py @@ -0,0 +1,196 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import RMSNorm +from fla.modules.feature_map import DPFPFeatureMap, HadamardFeatureMap, HedgehogFeatureMap, T2RFeatureMap +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class LinearAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: str = 1024, + expand_k: float = 1.0, + expand_v: float = 1.0, + num_heads: int = 8, + num_kv_heads: int | None = None, + feature_map: str = 'elementwise_product', + tie_feature_map_qk: bool = False, + output_norm: str = 'rmsnorm', + norm_q: bool = False, + norm_k: bool = False, + do_feature_map_norm: bool = False, + elementwise_affine: bool = True, + norm_eps: float = 1e-5, + layer_idx: int | None = None, + **kwargs, + ): + super().__init__() + + self.hidden_size = hidden_size + self.mode = mode + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + + assert mode in ['chunk', 'fused_chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + self.do_feature_map_norm = do_feature_map_norm + self.layer_idx = layer_idx + + if feature_map == 'hedgehog': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = HedgehogFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 't2r': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = T2RFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'elementwise_product': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = HadamardFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'dpfp': + self.feature_map_q = DPFPFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = DPFPFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'elu': + def elu(x): + return F.elu(x) + 1 + self.feature_map_q = elu + self.feature_map_k = elu + + elif feature_map == 'relu': + self.feature_map_q = nn.ReLU() + self.feature_map_k = nn.ReLU() + + elif feature_map == 'identity': + self.feature_map_q = nn.Identity() + self.feature_map_k = nn.Identity() + else: + raise NotImplementedError(f"Not supported feature map `{feature_map}`.") + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + + if output_norm == 'rmsnorm': + self.norm = RMSNorm(hidden_size=self.head_v_dim, elementwise_affine=elementwise_affine, + eps=norm_eps, dtype=torch.float32) + elif output_norm == 'identity': + self.norm = nn.Identity() + else: + raise NotImplementedError(f"Not supported output norm `{output_norm}`.") + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.norm_q = norm_q + self.norm_k = norm_k + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + # Match other recurrent layers: use the recurrent kernel for decode/small chunks. + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + last_state = get_layer_cache(self, past_key_values) + + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + if attention_mask is not None: + v = v.mul(attention_mask[:, -v.shape[-2]:, None]) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + if self.num_kv_groups > 1: + k = repeat(k, '... (h d) -> ... (h g) d', d=self.head_k_dim, g=self.num_kv_groups) + v = repeat(v, '... (h d) -> ... (h g) d', d=self.head_v_dim, g=self.num_kv_groups) + else: + k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + q = self.feature_map_q(q) + k = self.feature_map_k(k) + + if self.norm_q: + q = q / (q.sum(-1, True) + 1e-4) + if self.norm_k: + k = k / (k.sum(-1, True) + 1e-4) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, final_state = chunk_linear_attn( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + normalize=self.do_feature_map_norm, + ) + elif mode == 'fused_chunk': + o, final_state = fused_chunk_linear_attn( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + normalize=self.do_feature_map_norm, + ) + elif mode == 'fused_recurrent': + o, final_state = fused_recurrent_linear_attn( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + normalize=self.do_feature_map_norm, + ) + else: + raise NotImplementedError + update_layer_cache( + self, + past_key_values, + recurrent_state=final_state, + offset=q.shape[1], + ) + o = self.norm(o) + o = rearrange(o, '... h d -> ... (h d)') + o = self.o_proj(o) + return o, None, past_key_values diff --git a/fla/layers/log_linear_mamba2.py b/fla/layers/log_linear_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..20c2196c8dd4e50facd8000d5af4213ea2e31eb6 --- /dev/null +++ b/fla/layers/log_linear_mamba2.py @@ -0,0 +1,684 @@ +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.activations import ACT2FN +from transformers.utils import logging + +from fla.layers.mamba2 import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update, is_fast_path_available +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules.layernorm_gated import RMSNormGated, rmsnorm_fn +from fla.ops.log_linear_attn.chunk import LogLinearAttentionState, chunk_log_linear_attn + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +def ceil_log(x: int, b: int) -> int: + return math.ceil(math.log(x, b)) + + +def get_num_levels(length: int, base: int) -> int: + return ceil_log(length, base) + 1 + + +MAX_SEQUENCE_LENGTH = 2048 * 8 +LAMBDA_LEVEL_BASE = 2 +MAX_NUM_LEVELS = get_num_levels(length=MAX_SEQUENCE_LENGTH, base=LAMBDA_LEVEL_BASE) + + +def hmamba_chunk_scan_combined( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + dl: torch.Tensor, + L: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: LogLinearAttentionState | None = None, + seq_idx: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, +): + if z is not None: + raise NotImplementedError + if seq_idx is not None: + raise NotImplementedError + if cu_seqlens is not None: + raise NotImplementedError + if dt_softplus is not True: + raise NotImplementedError + if tuple(dt_limit) != (0.0, float("inf")): + raise NotImplementedError + if chunk_size != 64: + raise NotImplementedError + if not B.shape == C.shape: + raise ValueError("B and C must have the same shape") + + if D is not None: + if D.dim() != 1: + raise ValueError + D = rearrange(D, "h -> 1 1 h 1") + D_residual = x * D + + if dt_bias is not None: + dt = dt + rearrange(dt_bias, "h -> 1 1 h") + if dt_softplus: + dt = torch.nn.functional.softplus(dt) + if dt_limit != (0.0, float("inf")): + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + x = (x * rearrange(dt, "b l h -> b l h 1")).to(x.dtype) + A = rearrange(A, "h -> 1 1 h") * dt + + L = torch.nn.functional.softplus(rearrange(L, "h ell -> 1 1 h ell") * dl).to(L.dtype) + + y, state = chunk_log_linear_attn( + q=C, + k=B, + v=x, + g=A, + level_scales=L, + initial_state=initial_states, + output_final_state=return_final_states, + cu_seqlens=cu_seqlens, + ) + + if D is not None: + y = y + D_residual + + return y, state + + +def hmamba_split_conv1d_scan_combined( + zxbcdtdl: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor, + dt_bias: torch.Tensor, + A: torch.Tensor, + L: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + seq_idx: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + conv1d_fn=None, + conv_backend: str = "cuda", +) -> torch.Tensor: + """ + Argument: + zxbcdtdl: (batch, seqlen, 2 * dim + 2 * ngroups * dstate + nheads) where dim == nheads * headdim + conv1d_weight: (dim + 2 * ngroups * dstate, width) + conv1d_bias: (dim + 2 * ngroups * dstate,) + dt_bias: (nheads,) + A: (nheads) + L: (nheads, nlevels) + D: (nheads, headdim) or (nheads,) + initial_states: (batch, nheads, headdim, dstate) + seq_idx: (batch, seqlen), int32 + rmsnorm_weight: (dim,) + outproj_weight: (out_dim, dim) + outproj_bias: (out_dim,) + headdim: if D is 1D, headdim must be passed in + norm_before_gate: if True, we do RMSNorm(x) * F.silu(z). If False, we do RMSNorm(x * F.silu(z)) + Return: + out: (batch, seqlen, dim) + """ + if initial_states is not None: + raise NotImplementedError + if seq_idx is not None: + raise NotImplementedError + if dt_limit != (0.0, float("inf")): + raise NotImplementedError + if return_final_states is not False: + raise NotImplementedError + if norm_before_gate is not False: + raise NotImplementedError + if rmsnorm_weight is None: + raise NotImplementedError + if activation not in ["silu", "swish"]: + raise NotImplementedError + + batch, seqlen, _ = zxbcdtdl.shape + dlambda = L.shape[-1] + (nheads,) = D.shape + dim = nheads * headdim + dstate = (zxbcdtdl.shape[-1] - 2 * dim - nheads - nheads * dlambda) // ngroups // 2 + + if D.dim() != 1: + raise ValueError + if headdim is None: + raise ValueError + if nheads % ngroups != 0: + raise ValueError + if zxbcdtdl.shape != ( + batch, + seqlen, + 2 * dim + 2 * ngroups * dstate + nheads + nheads * dlambda, + ): + raise ValueError + if dt_bias.shape != (nheads,): + raise ValueError + if A.shape != (nheads,): + raise ValueError + if L.shape != (nheads, dlambda): + raise ValueError + if D.shape != (nheads,): + raise ValueError + if rmsnorm_weight is None: + raise ValueError + + zxBCdtl_splits = [dim, dim + 2 * ngroups * dstate, nheads, nheads * dlambda] + xBC_splits = [dim, ngroups * dstate, ngroups * dstate] + z, xBC, dt, dl = torch.split(zxbcdtdl, zxBCdtl_splits, dim=-1) + _conv_fn = conv1d_fn if conv1d_fn is not None else causal_conv1d_fn + _conv_out = _conv_fn( + rearrange(xBC, "b s d -> b d s"), + conv1d_weight, + bias=conv1d_bias, + activation=activation, + seq_idx=seq_idx, + ) + if conv_backend == 'triton': + _conv_out = _conv_out[0] + xBC = rearrange(_conv_out, "b d s -> b s d") + x, B, C = torch.split(xBC, xBC_splits, dim=-1) + x = rearrange(x, "b l (h p) -> b l h p", h=nheads, p=headdim) + B = rearrange(B, "b l (g n) -> b l g n", g=ngroups, n=dstate) + C = rearrange(C, "b l (g n) -> b l g n", g=ngroups, n=dstate) + dl = rearrange(dl, "b l (h ell) -> b l h ell", h=nheads, ell=dlambda) + y, _ = hmamba_chunk_scan_combined( + x=x, + dt=dt, + A=A, + B=B, + C=C, + dl=dl, + L=L, + chunk_size=chunk_size, + D=D, + z=z if rmsnorm_weight is None else None, + dt_bias=dt_bias, + dt_softplus=True, + seq_idx=seq_idx, + cu_seqlens=None, + dt_limit=dt_limit, + return_final_states=return_final_states, + ) + + y = rearrange(y, "b l h p -> b l (h p)") + if rmsnorm_weight is not None: + y = rmsnorm_fn( + x=y, + weight=rmsnorm_weight, + bias=None, + z=z, + eps=rmsnorm_eps, + group_size=None, + norm_before_gate=False, + ) + out = torch.nn.functional.linear(y, outproj_weight, outproj_bias) + return out + + +class LogLinearMamba2(nn.Module): + """ + Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. + A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) + ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, + and is why Mamba is called **selective** state spaces) + """ + + def __init__( + self, + num_heads: int, + head_dim: int = 64, + hidden_size: int = 2048, + state_size: int = 128, + expand: int = 2, + n_groups: int = 1, + conv_kernel: int = 4, + use_conv_bias: bool = False, + hidden_act: str = "silu", + rms_norm: bool = True, + chunk_size: int = 64, + time_step_rank: float = 256, + time_step_limit: tuple[float, float] = (0.0, float("inf")), + time_step_min: float = 0.001, + time_step_max: float = 0.1, + use_bias: bool = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + backend: str = "cuda", + ): + super().__init__() + self.num_heads = num_heads + self.hidden_size = hidden_size + self.ssm_state_size = state_size + self.conv_kernel_size = conv_kernel + self.intermediate_size = int(expand * self.hidden_size) + self.time_step_rank = int(time_step_rank) + self.layer_idx = layer_idx + self.use_conv_bias = use_conv_bias + self.activation = hidden_act + self.act = ACT2FN[hidden_act] + + self.layer_norm_epsilon = norm_eps + self.rms_norm = rms_norm + + self.n_groups = n_groups + self.head_dim = head_dim + self.chunk_size = chunk_size + + self.time_step_limit = time_step_limit + self.time_step_min = time_step_min + self.time_step_max = time_step_max + + self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=use_conv_bias, + kernel_size=conv_kernel, + groups=self.conv_dim, + padding=conv_kernel - 1, + ) + + self.num_lambda_dims = MAX_NUM_LEVELS + self.lambda_level_module = None + + # projection of the input hidden states + projection_size = ( + self.intermediate_size + + self.conv_dim + + self.num_heads * (self.num_lambda_dims + 1) + ) + self.in_proj = nn.Linear( + self.hidden_size, + projection_size, + bias=use_bias, + ) + # selective projection used to make dt, B and C input dependant + + # time step projection (discretization) + # instantiate once and copy inv_dt in init_weights of PretrainedModel + self.dt_bias = nn.Parameter(torch.ones(self.num_heads)) + + # S4D real initialization. These are not discretized! + # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded + A = torch.arange(1, self.num_heads + 1) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + + self.lambda_mode = "positive" + L = torch.ones(self.num_heads, self.num_lambda_dims) + self.L = nn.Parameter(L) + self.L._no_weight_decay = True + + self.norm = RMSNormGated( + self.intermediate_size, eps=self.layer_norm_epsilon, norm_before_gate=False, + ) + self.D = nn.Parameter(torch.ones(self.num_heads)) + self.D._no_weight_decay = True + + self.out_proj = nn.Linear( + self.intermediate_size, self.hidden_size, bias=use_bias, + ) + self.use_bias = use_bias + + if not is_fast_path_available: + logger.warning_once( + "The fast path is not available because one of " + "`(selective_state_update, causal_conv1d_fn, causal_conv1d_update)` is None. " + "Falling back to the naive implementation. " + "To install follow https://github.com/state-spaces/mamba/#installation and" + "https://github.com/Dao-AILab/causal-conv1d", + ) + import os + backend = os.environ.get('FLA_CONV_BACKEND', backend) + assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}" + if backend == 'cuda' and causal_conv1d_fn is None: + logger.warning_once( + "The CUDA backend is not available because `causal_conv1d` is None. " + "Falling back to the Triton backend. " + "To install follow https://github.com/Dao-AILab/causal-conv1d", + ) + backend = 'triton' + if backend == 'triton': + from fla.modules.convolution import causal_conv1d as causal_conv1d_triton + from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton + self.causal_conv1d_fn = causal_conv1d_triton + self.causal_conv1d_update = causal_conv1d_update_triton + logger.warning( + "LogLinearMamba2 does not recommend using Triton's conv1d backend, " + "as it is untested and may contain bugs.", + ) + else: + self.causal_conv1d_fn = causal_conv1d_fn + self.causal_conv1d_update = causal_conv1d_update + self.backend = backend + + def cuda_kernels_forward( + self, + hidden_states: torch.Tensor, + last_state: dict | None = None, + use_cache: bool = False, + attention_mask: torch.Tensor | None = None, + ): + if self.activation not in ["silu", "swish"]: + raise ValueError + + # 1. Gated MLP's linear projection + # Only apply padding mask during prefill (last_state is None). + # During decode, attention_mask has shape (B, accumulated_len) which + # mismatches hidden_states (B, 1, D). + hidden_states = apply_mask_to_padding_states( + hidden_states=hidden_states, + attention_mask=attention_mask if last_state is None else None, + ) + projected_states = self.in_proj(hidden_states) + + # Set up dimensions for reshapes later + batch_size, seq_len, _ = hidden_states.shape + groups_time_state_size = self.n_groups * self.ssm_state_size + d_mlp = ( + projected_states.shape[-1] + - 2 * self.intermediate_size + - 2 * self.n_groups * self.ssm_state_size + - self.num_heads * (self.num_lambda_dims + 1) + ) // 2 + if d_mlp != 0: + raise ValueError + + # Single step calculations via cache + if last_state is not None: + if hidden_states.shape[1] != 1: + raise ValueError("LogLinearMamba2 cached decoding only supports a single new token per step.") + + gate, xBC, dt, dl = torch.split( + projected_states.squeeze(1), + [ + self.intermediate_size, + self.conv_dim, + self.num_heads, + self.num_heads * self.num_lambda_dims, + ], + dim=-1, + ) + + # 2. Convolution sequence transformation + conv_state = last_state['conv_state'] + xBC = self.causal_conv1d_update( + xBC, + conv_state, + rearrange(self.conv1d.weight, "d 1 w -> d w"), + self.conv1d.bias, + self.activation, + ) + + x, B, C = torch.split( + xBC, + [ + self.intermediate_size, + groups_time_state_size, + groups_time_state_size, + ], + dim=-1, + ) + + # 3. SSM transformation + A = -torch.exp(self.A_log.float()) # (nheads,) + B = rearrange( + B, + "b (g n) -> b g n", + b=batch_size, + g=self.n_groups, + n=self.ssm_state_size, + ) + C = rearrange( + C, + "b (g n) -> b g n", + b=batch_size, + g=self.n_groups, + n=self.ssm_state_size, + ) + x_reshaped = rearrange( + x, + "b (h p) -> b h p", + b=batch_size, + h=self.num_heads, + p=self.head_dim, + ) + dl_reshaped = rearrange( + dl, + "b (h ell) -> b h ell", + b=batch_size, + h=self.num_heads, + ell=self.num_lambda_dims, + ) + y, hssm_state = hmamba_chunk_scan_combined( + x_reshaped, + dt=dt, + A=A, + B=B, + C=C, + dl=dl_reshaped, + L=self.L, + D=self.D, + z=None, + dt_bias=self.dt_bias, + dt_softplus=True, + initial_states=last_state['recurrent_state'], + return_final_states=True, + ) + y = rearrange( + y, + "b h p -> b (h p)", + b=batch_size, + h=self.num_heads, + p=self.head_dim, + ) + y = self.norm(y, gate) + + # 4. Final linear projection + out = self.out_proj(y)[:, None, ...] + return out, conv_state, hssm_state + + # Fused calculations or step by step if no initialized cache is found + else: + A = -torch.exp( + self.A_log.float(), + ) # (num_heads) or (intermediate_size, state_size) + dt_limit_kwargs = ( + {} + if self.time_step_limit == (0.0, float("inf")) + else {"dt_limit": self.time_step_limit} + ) + + # 2-4. Fused kernel for conv1d, SSM, and the final projection + if self.training and not use_cache: + out = torch.utils.checkpoint.checkpoint( + hmamba_split_conv1d_scan_combined, + use_reentrant=False, + # function arguments + zxbcdtdl=projected_states, + conv1d_weight=rearrange(self.conv1d.weight, "d 1 w -> d w"), + conv1d_bias=self.conv1d.bias, + dt_bias=self.dt_bias, + A=A, + L=self.L, + D=self.D, + chunk_size=self.chunk_size, + conv1d_fn=self.causal_conv1d_fn, + conv_backend=self.backend, + seq_idx=None, # was seq_idx + activation=self.activation, + rmsnorm_weight=self.norm.weight, + rmsnorm_eps=self.norm.eps, + outproj_weight=self.out_proj.weight, + outproj_bias=self.out_proj.bias, + headdim=self.head_dim, + ngroups=self.n_groups, + norm_before_gate=False, + return_final_states=False, + **dt_limit_kwargs, + ) + return out, None, None + + else: + gate, xBC, dt, dl = torch.split( + projected_states, + [ + self.intermediate_size, + self.conv_dim, + self.num_heads, + self.num_heads * self.num_lambda_dims, + ], + dim=-1, + ) + + # 2. Convolution sequence transformation + # Init cache + masked_xBC = apply_mask_to_padding_states(xBC, attention_mask) + new_conv_state = None + if use_cache: + xBC_t = rearrange(masked_xBC, "b l d -> b d l") + new_conv_state = torch.nn.functional.pad( + xBC_t, + (self.conv_kernel_size - xBC_t.shape[-1], 0), + ) + + _conv1d_output = self.causal_conv1d_fn( + x=xBC.transpose(1, 2), + weight=rearrange(self.conv1d.weight, "d 1 w -> d w"), + bias=self.conv1d.bias, + activation=self.activation, + ) + if self.backend == 'cuda': + xBC = _conv1d_output.transpose(1, 2) + elif self.backend == 'triton': + xBC = _conv1d_output[0].transpose(1, 2).contiguous() + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + xBC = apply_mask_to_padding_states( + hidden_states=xBC, + attention_mask=attention_mask, + ) + + x, B, C = torch.split( + xBC, + [ + self.intermediate_size, + groups_time_state_size, + groups_time_state_size, + ], + dim=-1, + ) + + # 3. SSM transformation + y, hssm_state = hmamba_chunk_scan_combined( + rearrange( + x, + "b l (h p) -> b l h p", + b=batch_size, + l=seq_len, + p=self.head_dim, + ), + dt=dt, + A=A, + B=rearrange( + B, + "b l (g n) -> b l g n", + b=batch_size, + l=seq_len, + g=self.n_groups, + ), + C=rearrange( + C, + "b l (g n) -> b l g n", + b=batch_size, + l=seq_len, + g=self.n_groups, + ), + dl=rearrange( + dl, + "b l (h ell) -> b l h ell", + b=batch_size, + h=self.num_heads, + ell=self.num_lambda_dims, + ), + L=self.L, + chunk_size=self.chunk_size, + D=self.D, + z=None, + seq_idx=None, + return_final_states=True, + dt_bias=self.dt_bias, + dt_softplus=True, + **dt_limit_kwargs, + ) + + y = rearrange( + y, + "b l h p -> b l (h p)", + b=batch_size, + l=seq_len, + h=self.num_heads, + p=self.head_dim, + ) + # Multiply "gate" branch and apply extra normalization layer + y = self.norm(y, gate) + + # 4. Final linear projection + out = self.out_proj(y) + + return out, new_conv_state, hssm_state + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + last_state = get_layer_cache(self, past_key_values) + + if "cuda" in self.in_proj.weight.device.type: + output, conv_state, hssm_state = self.cuda_kernels_forward( + hidden_states, last_state, use_cache, attention_mask + ) + else: + raise NotImplementedError + + update_layer_cache( + self, + past_key_values, + recurrent_state=hssm_state, + conv_state=conv_state, + offset=hidden_states.shape[1], + ) + + return output, None, past_key_values diff --git a/fla/layers/mamba.py b/fla/layers/mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..b9d08b95c6436021f10bc1703bb44fee7bfbcea9 --- /dev/null +++ b/fla/layers/mamba.py @@ -0,0 +1,395 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from transformers.utils import logging + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules.activations import ACT2FN + +with warnings.catch_warnings(): + warnings.simplefilter('ignore') + try: + from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn + from mamba_ssm.ops.triton.selective_state_update import selective_state_update + except ImportError: + selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None + + try: + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + except ImportError: + causal_conv1d_update, causal_conv1d_fn = None, None + is_fast_path_available = all(( + selective_state_update, + selective_scan_fn, + mamba_inner_fn, + )) +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class Mamba(nn.Module): + """ + Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. + A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) + ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, + and is why Mamba is called **selective** state spaces) + """ + + def __init__( + self, + hidden_size: int = 2048, + state_size: int = 16, + conv_kernel: int = 4, + use_conv_bias: bool = True, + intermediate_size: int = 2048, + time_step_rank: int = 256, + use_bias: bool = True, + hidden_act: str = "silu", + layer_idx: int = None, + backend: str = "cuda", + ): + super().__init__() + + self.hidden_size = hidden_size + self.ssm_state_size = state_size + self.conv_kernel_size = conv_kernel + self.use_conv_bias = use_conv_bias + self.intermediate_size = intermediate_size + self.time_step_rank = time_step_rank + self.use_bias = use_bias + + self.conv1d = nn.Conv1d( + in_channels=self.intermediate_size, + out_channels=self.intermediate_size, + bias=use_conv_bias, + kernel_size=conv_kernel, + groups=self.intermediate_size, + padding=conv_kernel - 1, + ) + + self.activation = hidden_act + self.act = ACT2FN[hidden_act] + + self.layer_idx = layer_idx + + # projection of the input hidden states + self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=use_bias) + # selective projection used to make dt, B and C input dependant + self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False) + # time step projection (discretization) + self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True) + + # S4D real initialization. These are not discretized! + # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded + A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32)[None, :] + A = A.expand(self.intermediate_size, -1).contiguous() + + self.A_log = nn.Parameter(torch.log(A)) + self.D = nn.Parameter(torch.ones(self.intermediate_size)) + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=use_bias) + + if not is_fast_path_available: + logger.warning_once( + "The fast path is not available because on of " + "`(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" + " is None. Falling back to the naive implementation. " + "To install follow https://github.com/state-spaces/mamba/#installation and" + " https://github.com/Dao-AILab/causal-conv1d", + ) + import os + backend = os.environ.get('FLA_CONV_BACKEND', backend) + assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}" + if backend == 'cuda' and causal_conv1d_fn is None: + logger.warning_once( + "The CUDA backend is not available because `causal_conv1d` is None. " + "Falling back to the Triton backend. " + "To install follow https://github.com/Dao-AILab/causal-conv1d", + ) + backend = 'triton' + if backend == 'triton': + from fla.modules.convolution import causal_conv1d as causal_conv1d_triton + from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton + self.causal_conv1d_fn = causal_conv1d_triton + self.causal_conv1d_update = causal_conv1d_update_triton + else: + self.causal_conv1d_fn = causal_conv1d_fn + self.causal_conv1d_update = causal_conv1d_update + self.backend = backend + + def _to_causal_conv_layout(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states.transpose(1, 2).contiguous() + + def _from_causal_conv_layout(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states.transpose(1, 2).contiguous() + + def _build_conv_state(self, hidden_states: torch.Tensor) -> torch.Tensor: + seq_len = hidden_states.shape[-1] + if seq_len >= self.conv_kernel_size: + return hidden_states[..., -self.conv_kernel_size:].contiguous() + return nn.functional.pad(hidden_states, (self.conv_kernel_size - seq_len, 0)).contiguous() + + def cuda_kernels_forward( + self, + hidden_states: torch.Tensor, + last_state: dict | None = None, + use_cache: bool | None = False, + attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[dict], + ): + if last_state is not None and hidden_states.shape[1] != 1: + raise ValueError("Mamba cached decoding only supports a single new token per step.") + + # 1. Gated MLP's linear projection + projected_states = self.in_proj(hidden_states).transpose(1, 2) + + if self.training and not use_cache: + contextualized_states = mamba_inner_fn( + projected_states, + self.conv1d.weight, + self.conv1d.bias if self.use_conv_bias else None, + self.x_proj.weight, + self.dt_proj.weight, + self.out_proj.weight, + self.out_proj.bias.float() if self.use_bias else None, + -torch.exp(self.A_log.float()), + None, # input-dependent B + None, # input-dependent C + self.D.float(), + delta_bias=self.dt_proj.bias.float(), + delta_softplus=True, + ) + return contextualized_states, None, None + + hidden_states, gate = projected_states.chunk(2, dim=1) + + if attention_mask is not None and last_state is None: + # Mask before the depthwise conv so cached/prefill conv inputs do not keep pad tokens. + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 2. Convolution sequence transformation + conv_inputs = hidden_states + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) + if last_state is not None: + conv_state = last_state['conv_state'] + ssm_state = last_state['recurrent_state'] + + if self.backend == 'triton': + hidden_states, conv_state = self.causal_conv1d_update( + x=self._to_causal_conv_layout(conv_inputs), + cache=conv_state, + weight=conv_weights, + bias=self.conv1d.bias, + activation=self.activation, + ) + hidden_states = self._from_causal_conv_layout(hidden_states) + else: + hidden_states = self.causal_conv1d_update( + conv_inputs.squeeze(-1), + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + ) + hidden_states = hidden_states.unsqueeze(-1) + else: + conv_state = None + ssm_state = None + if self.backend == 'triton': + hidden_states, conv_state = self.causal_conv1d_fn( + x=self._to_causal_conv_layout(conv_inputs), + weight=conv_weights, + bias=self.conv1d.bias, + activation=self.activation, + output_final_state=bool(use_cache), + ) + hidden_states = self._from_causal_conv_layout(hidden_states) + else: + if use_cache: + conv_state = self._build_conv_state(conv_inputs) + hidden_states = self.causal_conv1d_fn( + conv_inputs, conv_weights, self.conv1d.bias, activation=self.activation, + ) + + if attention_mask is not None and last_state is None: + # Re-mask after the conv: causal kernels can regenerate non-zero values at masked positions, + # and those values would otherwise leak into x_proj and the SSM recurrence. + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 3. State Space Model sequence transformation + # 3.a. input varying initialization of time_step, B and C + ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) + time_step, B, C = torch.split( + ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1, + ) + discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + + A = -torch.exp(self.A_log.float()) + # 3.c perform the recurrence y ← SSM(A, B, C)(x) + time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None + if last_state is not None: + scan_outputs = selective_state_update( + ssm_state, + hidden_states[..., 0], + discrete_time_step[..., 0], + A, + B[:, 0], + C[:, 0], + self.D, + gate[..., 0], + time_proj_bias, + dt_softplus=True, + ).unsqueeze(-1) + else: + scan_outputs, ssm_state = selective_scan_fn( + hidden_states, + discrete_time_step, + A, + B.transpose(1, 2), + C.transpose(1, 2), + self.D.float(), + gate, + time_proj_bias, + delta_softplus=True, + return_last_state=True, + ) + + # 4. Final linear projection + contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) + return contextualized_states, conv_state, ssm_state + + def slow_forward( + self, + input_states, + last_state: dict | None = None, + use_cache: bool | None = False, + attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[dict], + ): + if last_state is not None and input_states.shape[1] != 1: + raise ValueError("Mamba cached decoding only supports a single new token per step.") + + batch_size, seq_len, _ = input_states.shape + dtype = input_states.dtype + # 1. Gated MLP's linear projection + # [batch, 2 * intermediate_size, seq_len] + projected_states = self.in_proj(input_states).transpose(1, 2) + hidden_states, gate = projected_states.chunk(2, dim=1) + + if attention_mask is not None and last_state is None: + # Mask before the depthwise conv so cached/prefill conv inputs do not keep pad tokens. + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 2. Convolution sequence transformation + if last_state is not None: + conv_state = last_state['conv_state'] + ssm_state = last_state['recurrent_state'].clone().to(hidden_states.device) + + # decode path: single token + conv_state = conv_state.roll(shifts=-1, dims=-1) + conv_state[:, :, -1] = hidden_states[:, :, 0].to(conv_state.device) + hidden_states = torch.sum(conv_state * self.conv1d.weight[:, 0, :], dim=-1) + if self.use_conv_bias: + hidden_states += self.conv1d.bias + # [batch, intermediate_size, 1] : decoding + hidden_states = self.act(hidden_states).to(dtype).unsqueeze(-1) + elif use_cache: + ssm_state = torch.zeros( + (batch_size, self.intermediate_size, self.ssm_state_size), + device=hidden_states.device, dtype=dtype, + ) + conv_state = self._build_conv_state(hidden_states) + # [batch, intermediate_size, seq_len] + hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) + else: + ssm_state = torch.zeros( + (batch_size, self.intermediate_size, self.ssm_state_size), + device=hidden_states.device, dtype=dtype, + ) + conv_state = None + # [batch, intermediate_size, seq_len] + hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) + + if attention_mask is not None and last_state is None: + # Re-mask after the conv: causal kernels can regenerate non-zero values at masked positions, + # and those values would otherwise leak into x_proj and the SSM recurrence. + hidden_states = hidden_states * attention_mask.unsqueeze(1) + + # 3. State Space Model sequence transformation + # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] + ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) + time_step, B, C = torch.split( + ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1, + ) + # [batch, seq_len, intermediate_size] + discrete_time_step = self.dt_proj(time_step) + # [batch, intermediate_size, seq_len] + discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) + + # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) + # [intermediate_size, ssm_state_size] + A = -torch.exp(self.A_log.float()) + # [batch, intermediate_size, seq_len, ssm_state_size] + discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) + # [batch, intermediate_size, seq_len, ssm_state_size] + discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + # 3.c perform the recurrence y ← SSM(A, B, C)(x) + scan_outputs = [] + for i in range(hidden_states.shape[-1]): + # [batch, intermediade_size, ssm_state] + ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] + # [batch, intermediade_size, 1] + scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) + scan_outputs.append(scan_output[:, :, 0]) + # [batch, seq_len, intermediade_size] + scan_output = torch.stack(scan_outputs, dim=-1) + scan_output = scan_output + (hidden_states * self.D[None, :, None]) + scan_output = (scan_output * self.act(gate)) + + # 4. Final linear projection + # [batch, seq_len, hidden_size] + contextualized_states = self.out_proj(scan_output.transpose(1, 2)) + return contextualized_states, conv_state, ssm_state + # fmt: on + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + last_state = get_layer_cache(self, past_key_values) + + if is_fast_path_available and "cuda" in self.x_proj.weight.device.type: + output, conv_state, ssm_state = self.cuda_kernels_forward( + hidden_states, last_state, use_cache, attention_mask, **kwargs + ) + else: + output, conv_state, ssm_state = self.slow_forward( + hidden_states, last_state, use_cache, attention_mask, **kwargs + ) + + if use_cache and past_key_values is not None: + update_layer_cache( + self, + past_key_values, + recurrent_state=ssm_state, + conv_state=conv_state, + offset=hidden_states.shape[1], + ) + + return output, None, past_key_values diff --git a/fla/layers/mamba2.py b/fla/layers/mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..397bdb34b88e927abefa6e23eb971d199a66c5fd --- /dev/null +++ b/fla/layers/mamba2.py @@ -0,0 +1,649 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from transformers.utils import logging + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules.activations import ACT2FN +from fla.modules.layernorm_gated import RMSNormGated + +with warnings.catch_warnings(): + warnings.simplefilter('ignore') + try: + from mamba_ssm.ops.triton.selective_state_update import selective_state_update + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined + except ImportError: + selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined = None, None, None + try: + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + except ImportError: + causal_conv1d_update, causal_conv1d_fn = None, None + is_fast_path_available = selective_state_update is not None + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +def apply_mask_to_padding_states(hidden_states, attention_mask): + """ + Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 + """ + if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + dtype = hidden_states.dtype + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) + + return hidden_states + + +def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int): + """ + Padding x tensor with `pad_size` on the seq_len dim (dim=1) + + Assumes that we only have tensors of either size 4 or 3 + """ + pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0) + + return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0) + + +def reshape_into_chunks(input_tensor, pad_size, chunk_size): + """ + Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and + simultaneously splitting it into chunk sequences. + + Assumes that we only have tensors of either size 4 or 3 + """ + # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...] + input_tensor = pad_tensor_by_size(input_tensor, pad_size) + + if len(input_tensor.shape) == 3: + # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads] + return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2]) + else: + # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> + # [bsz, -1, chunk_size, num_heads, head_dim or state_size] + return input_tensor.reshape( + input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3], + ) + + +def segment_sum(input_tensor): + """ + More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions. + """ + chunk_size = input_tensor.size(-1) + # 1. expand input tensor to have an additional dimension and repeat along that dimension + # [..., chunk_size] -> [..., chunk_size, chunk_size] + input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size) + # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag + mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1) + input_tensor = input_tensor.masked_fill(~mask, 0) + # 3. compute actual cumsum + tensor_segsum = torch.cumsum(input_tensor, dim=-2) + + # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time) + mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0) + tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf) + return tensor_segsum + + +class Mamba2(nn.Module): + """ + Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. + A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) + ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, + and is why Mamba is called **selective** state spaces) + """ + + def __init__( + self, + num_heads: int, + head_dim: int = 64, + hidden_size: int = 2048, + state_size: int = 128, + expand: int = 2, + n_groups: int = 1, + conv_kernel: int = 4, + use_conv_bias: bool = False, + hidden_act: str = "silu", + rms_norm: bool = True, + chunk_size: int = 256, + time_step_rank: float = 256, + time_step_limit: tuple[float, float] = (0.0, float("inf")), + time_step_min: float = 0.001, + time_step_max: float = 0.1, + use_bias: bool = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + backend: str = "cuda", + ) -> Mamba2: + super().__init__() + + self.num_heads = num_heads + self.head_dim = head_dim + self.hidden_size = hidden_size + self.ssm_state_size = state_size + self.expand = expand + self.intermediate_size = int(expand * hidden_size) + self.n_groups = n_groups + + self.conv_kernel_size = conv_kernel + self.use_conv_bias = use_conv_bias + self.activation = hidden_act + self.act = ACT2FN[hidden_act] + + self.rms_norm = rms_norm + self.norm_eps = norm_eps + + self.chunk_size = chunk_size + + self.time_step_rank = int(time_step_rank) + self.time_step_limit = time_step_limit + self.time_step_min = time_step_min + self.time_step_max = time_step_max + + self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=use_conv_bias, + kernel_size=conv_kernel, + groups=self.conv_dim, + padding=conv_kernel - 1, + ) + + # projection of the input hidden states + projection_size = self.intermediate_size + self.conv_dim + self.num_heads + self.in_proj = nn.Linear( + self.hidden_size, + projection_size, + bias=use_bias, + ) + # selective projection used to make dt, B and C input dependant + + # time step projection (discretization) + # instantiate once and copy inv_dt in init_weights of PretrainedModel + # hard coded for now + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_heads) * ( + math.log(self.time_step_max) - math.log(self.time_step_min) + ) + math.log(self.time_step_min) + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + + # S4D real initialization. These are not discretized! + # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded + A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + self.norm = RMSNormGated( + self.intermediate_size, eps=self.norm_eps, norm_before_gate=False, + ) + self.D = nn.Parameter(torch.ones(self.num_heads)) + self.D._no_weight_decay = True + + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=use_bias) + self.use_bias = use_bias + + self.layer_idx = layer_idx + + if not is_fast_path_available: + logger.warning_once( + "The fast path is not available because one of " + "`(selective_state_update)` is None. " + "Falling back to the naive implementation. " + "To install follow https://github.com/state-spaces/mamba/#installation", + ) + import os + backend = os.environ.get('FLA_CONV_BACKEND', backend) + assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}" + if backend == 'cuda' and causal_conv1d_fn is None: + logger.warning_once( + "The CUDA backend is not available because `causal_conv1d` is None. " + "Falling back to the Triton backend. " + "To install follow https://github.com/Dao-AILab/causal-conv1d", + ) + backend = 'triton' + if backend == 'triton': + from fla.modules.convolution import causal_conv1d as causal_conv1d_triton + from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton + self.causal_conv1d_fn = causal_conv1d_triton + self.causal_conv1d_update = causal_conv1d_update_triton + logger.warning( + "Mamba2 does not recommend using Triton's conv1d backend, " + "as it is untested and may contain bugs.", + ) + else: + self.causal_conv1d_fn = causal_conv1d_fn + self.causal_conv1d_update = causal_conv1d_update + self.backend = backend + + def cuda_kernels_forward( + self, + hidden_states: torch.Tensor, + last_state: dict | None = None, + use_cache: bool = False, + attention_mask: torch.Tensor | None = None, + ): + # 1. Gated MLP's linear projection + projected_states = self.in_proj(hidden_states) + + # Set up dimensions for reshapes later + batch_size, seq_len, _ = hidden_states.shape + groups_time_state_size = self.n_groups * self.ssm_state_size + d_mlp = ( + projected_states.shape[-1] + - 2 * self.intermediate_size + - 2 * self.n_groups * self.ssm_state_size + - self.num_heads + ) // 2 + + # Single step calculations via cache (decode) + if last_state is not None: + if hidden_states.shape[1] != 1: + raise ValueError("Mamba2 cached decoding only supports a single new token per step.") + conv_state = last_state['conv_state'] + ssm_state = last_state['recurrent_state'] + + _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split( + [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1, + ) + + # 2. Convolution sequence transformation + hidden_states_B_C = self.causal_conv1d_update( + hidden_states_B_C.contiguous(), + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + self.activation, + ) + + hidden_states, B, C = torch.split( + hidden_states_B_C, + [ + self.intermediate_size, + groups_time_state_size, + groups_time_state_size, + ], + dim=-1, + ) + + # 3. SSM transformation + A = -torch.exp(self.A_log.float()) # (nheads,) + A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + dt = dt[:, :, None].expand(-1, -1, self.head_dim) + dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) + D = self.D[:, None, ...].expand(-1, self.head_dim) + B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups) + C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups) + hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) + + hidden_states = selective_state_update( + ssm_state, + hidden_states_reshaped, + dt, + A, + B, + C, + D, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + ) + hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim) + hidden_states = self.norm(hidden_states, gate) + + # 4. Final linear projection + out = self.out_proj(hidden_states)[:, None, ...] + + # conv_state is updated in-place by causal_conv1d_update + # ssm_state is updated in-place by selective_state_update + return out, conv_state, ssm_state + + # Fused calculations or step by step if no initialized cache is found (prefill) + else: + A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) + dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} + + # 2-4. Fused kernel for conv1d, SSM, and the final projection + if self.training and not use_cache: + out = mamba_split_conv1d_scan_combined( + projected_states, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + self.dt_bias, + A, + D=self.D, + chunk_size=self.chunk_size, + seq_idx=None, # was seq_idx + activation=self.activation, + rmsnorm_weight=self.norm.weight, + rmsnorm_eps=self.norm.eps, + outproj_weight=self.out_proj.weight, + outproj_bias=self.out_proj.bias, + headdim=self.head_dim, + ngroups=self.n_groups, + norm_before_gate=False, + return_final_states=False, + **dt_limit_kwargs, + ) + return out, None, None + + else: + _, _, gate, hidden_states_B_C, dt = projected_states.split( + [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1, + ) + + # 2. Convolution sequence transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + # Compute conv_state for cache + new_conv_state = None + if use_cache: + hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2) + new_conv_state = nn.functional.pad( + hidden_states_B_C_transposed, + (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0), + ) + + if self.activation not in ["silu", "swish"]: + hidden_states_B_C = self.act( + self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2), + ) + else: + _conv1d_output = self.causal_conv1d_fn( + x=hidden_states_B_C.transpose(1, 2).contiguous(), + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + ) + if self.backend == 'cuda': + hidden_states_B_C = _conv1d_output + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + elif self.backend == 'triton': + hidden_states_B_C, _ = _conv1d_output + hidden_states_B_C = hidden_states_B_C.transpose(1, 2).contiguous() + else: + raise ValueError(f"Unsupported backend: {self.backend}") + + hidden_states_B_C = (hidden_states_B_C * attention_mask[:, :, None]).to(hidden_states_B_C.dtype) \ + if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1 \ + else hidden_states_B_C + hidden_states, B, C = torch.split( + hidden_states_B_C, + [self.intermediate_size, groups_time_state_size, groups_time_state_size], + dim=-1, + ) + + # 3. SSM transformation + scan_output, ssm_state = mamba_chunk_scan_combined( + hidden_states.view(batch_size, seq_len, -1, self.head_dim), + dt, + A, + B.view(batch_size, seq_len, self.n_groups, -1), + C.view(batch_size, seq_len, self.n_groups, -1), + chunk_size=self.chunk_size, + D=self.D, + z=None, + seq_idx=None, + return_final_states=True, + dt_bias=self.dt_bias, + dt_softplus=True, + **dt_limit_kwargs, + ) + + scan_output = scan_output.view(batch_size, seq_len, -1) + # Multiply "gate" branch and apply extra normalization layer + scan_output = self.norm(scan_output, gate) + + # 4. Final linear projection + out = self.out_proj(scan_output) + + return out, new_conv_state, ssm_state + + # fmt: off + def torch_forward( + self, + input_states, + last_state: dict | None = None, + use_cache: bool = False, + attention_mask: torch.Tensor | None = None, + ): + batch_size, seq_len, _ = input_states.shape + dtype = input_states.dtype + + # 1. Gated MLP's linear projection + projected_states = self.in_proj(input_states) + d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - + 2 * self.n_groups * self.ssm_state_size - self.num_heads) // 2 + _, _, gate, hidden_states_B_C, dt = projected_states.split( + [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1, + ) + + # 2. Convolution sequence transformation + if last_state is not None: + if input_states.shape[1] != 1: + raise ValueError("Mamba2 cached decoding only supports a single new token per step.") + # Decode path: single-step update + conv_state = last_state['conv_state'] + ssm_state = last_state['recurrent_state'] + + conv_state = conv_state.roll(shifts=-1, dims=-1) + conv_state[:, :, -1] = hidden_states_B_C[:, 0, :].to(conv_state.device) + + # We need to guarantee that anything regarding the cache is on the same device + conv_states_for_compute = conv_state.to(device=self.conv1d.weight.device) + + hidden_states_B_C = torch.sum( + conv_states_for_compute * self.conv1d.weight.squeeze(1), dim=-1, + ) + if self.use_conv_bias: + hidden_states_B_C = hidden_states_B_C + self.conv1d.bias + hidden_states_B_C = self.act(hidden_states_B_C) + else: + # Prefill path + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + new_conv_state = None + if use_cache: + hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2) + new_conv_state = nn.functional.pad( + hidden_states_B_C_transposed, (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0), + ) + + hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)) + + if last_state is None: + hidden_states_B_C = (hidden_states_B_C * attention_mask[:, :, None]).to(hidden_states_B_C.dtype) \ + if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1 \ + else hidden_states_B_C + hidden_states, B, C = torch.split( + hidden_states_B_C, + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], + dim=-1, + ) + + # 3. SSM transformation + A = -torch.exp(self.A_log.float()) # [num_heads] + if last_state is not None: + # Decode path + cache_device = ssm_state.device + + # Note: there is no need to pad parameter matrices here, as there is just one new token + # for batched generation + dt = dt[:, 0, :][:, None, ...] + dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) + # [num_heads] -> [num_heads, head_dim] + dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) + + dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype)) + dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) + A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + # [bsz, num_heads, head_dim, state_size] + dA = (torch.exp(dt[..., None] * A)).to(device=cache_device) + + # Discretize B + # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] -> + # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size] + B = B.reshape(batch_size, self.n_groups, -1)[..., None, :] + B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() + B = B.reshape(batch_size, -1, B.shape[-1]) + # [bsz, num_heads, head_dim, state_size] + dB = dt[..., None] * B[..., None, :] + + # Discretize x into dB + # [bsz, intermediate_size] -> [bsz, num_heads, head_dim] + hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) + dBx = (dB * hidden_states[..., None]).to(device=cache_device) + + # State calculation + ssm_state = ssm_state * dA + dBx + + # Subsequent output + # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size] + C = C.reshape(batch_size, self.n_groups, -1)[..., None, :] + C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() + C = C.reshape(batch_size, -1, C.shape[-1]) + # [bsz, num_heads, head_dim] + + ssm_states_for_compute = ssm_state.to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n] + # Reshape ssm_states to merge the first two dimensions + # Shape: [b*h, d, n] + ssm_states_reshaped = ssm_states_for_compute.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) + C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1] + y = torch.bmm(ssm_states_reshaped, C_reshaped) + y = y.view(batch_size, self.num_heads, self.head_dim) + + # D skip connection + # [num_heads] -> [num_heads, head_dim] + D = self.D[..., None].expand(self.D.shape[0], self.head_dim) + y = (y + hidden_states * D).to(y.dtype) + + # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size] + y = y.reshape(batch_size, -1)[:, None, ...] + + scan_output = self.norm(y, gate) + contextualized_states = self.out_proj(scan_output.to(dtype)) + return contextualized_states, conv_state, ssm_state + else: + # Prefill path + # begin ssd naive implementation without einsums + dt = nn.functional.softplus(dt + self.dt_bias) + dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) + hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() + B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() + C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() + B = B.repeat(1, 1, self.num_heads // self.n_groups, 1) + C = C.repeat(1, 1, self.num_heads // self.n_groups, 1) + pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size + + D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None] + A = A.to(hidden_states.dtype) * dt + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)] + + # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size] + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + # shape: (b, c, l, s, h, n) + G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] + G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h) + + # Compute M, equivalent to applying attention mask to weights + M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] + M = M_intermediate.sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = torch.zeros_like(states[:, :1]) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) + decay_chunk = decay_chunk.transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, ssm_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = (C[..., None, :] * states[:, :, None, ...]) + state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) + Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None]) + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + y = Y_diag + Y_off + # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim] + y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) + + y = y + D_residual + # Cutting off padded chunks + if pad_size > 0: + y = y[:, :seq_len, :, :] + y = y.reshape(batch_size, seq_len, -1) + + scan_output = self.norm(y, gate) + + # end ssd naive + + # 4. Final linear projection + contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size] + return contextualized_states, new_conv_state if use_cache else None, ssm_state + # fmt: on + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + last_state = get_layer_cache(self, past_key_values) + + if is_fast_path_available and "cuda" in self.in_proj.weight.device.type: + output, conv_state, ssm_state = self.cuda_kernels_forward(hidden_states, last_state, use_cache, attention_mask) + else: + dtype = hidden_states.dtype + if last_state is None and attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) + output, conv_state, ssm_state = self.torch_forward(hidden_states, last_state, use_cache, attention_mask) + + update_layer_cache( + self, + past_key_values, + recurrent_state=ssm_state, + conv_state=conv_state, + offset=hidden_states.shape[1], + ) + + return output, None, past_key_values diff --git a/fla/layers/mesa_net.py b/fla/layers/mesa_net.py new file mode 100644 index 0000000000000000000000000000000000000000..5aae752f9b3c280559fbbae9f84d260496bb4634 --- /dev/null +++ b/fla/layers/mesa_net.py @@ -0,0 +1,221 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.l2norm import l2_norm +from fla.ops.mesa_net import chunk_mesa_net, mesa_net_decoding_one_step + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class MesaNet(nn.Module): + """ + The layer implementaion for [MesaNet: Sequence Modeling by Locally Optimal Test-Time Training]. # noqa + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1. + num_heads (int, Optional): + The number of heads. Default: 16. + mode (str, Optional): + Which MesaNet kernel to use. + Currently available: `chunk`. + Default: `chunk`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `False`. + conv_size (int): + The kernel size of the short convolution. Default: 4. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + lambda_lower_bound (float): + The lower bound for the lambda parameter. Default: 0.25. + max_cg_step_training (int): + The maximum number of CG steps for training. Default: 30. + max_cg_step_decoding (int): + The maximum number of CG steps for decoding. Default: 30. + """ + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 16, + head_dim: int = 128, + mode: str = 'chunk', + use_output_gate: bool = False, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + lambda_lower_bound: float = 0.25, + max_cg_step_training: int = 30, + max_cg_step_decoding: int = 30, + **kwargs, + ) -> MesaNet: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.num_heads = num_heads + self.head_dim = head_dim + self.key_dim = self.num_heads * self.head_dim + self.value_dim = self.key_dim + self.head_k_dim = self.head_dim + self.head_v_dim = self.head_dim + self.layer_idx = layer_idx + self.lambda_lower_bound = lambda_lower_bound + self.max_cg_step_training = max_cg_step_training + self.max_cg_step_decoding = max_cg_step_decoding + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=True) + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=True) + + lambda_initial_value = 1.0 + init_lamb_value = torch.log(torch.exp(torch.tensor(lambda_initial_value - lambda_lower_bound)) - 1.0) + init_lamb_params = torch.empty(self.key_dim, dtype=torch.float32).fill_(init_lamb_value) + + self.lambda_params = nn.Parameter(init_lamb_params) + self.lambda_params._no_weight_decay = True + + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=self.conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=self.conv_bias, + activation='silu', + ) + if use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + conv_state_q, conv_state_k = None, None + if last_state is not None: + conv_state_q, conv_state_k = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v = self.v_proj(hidden_states) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + beta = self.b_proj(hidden_states).float().sigmoid() + g = F.logsigmoid(self.a_proj(hidden_states).float()) + lamb = F.softplus(self.lambda_params.float()) + self.lambda_lower_bound + lamb = lamb.reshape(self.num_heads, -1) + + last_h_kk, last_h_kv = last_state['recurrent_state'] if last_state is not None else (None, None) + + # prefilling or training + # Note that QK will be normalized inside the kernel to avoid saving the activations, thereby reducing the memory usage. + if last_state is None: + o, h_kk, h_kv = chunk_mesa_net( + q=q, + k=k, + v=v, + g=g, + beta=beta, + lamb=lamb, + output_final_state=use_cache, + max_CG_iteration=self.max_cg_step_training, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + # decoding + else: + q = l2_norm(q) + k = l2_norm(k) + o, h_kk, h_kv = mesa_net_decoding_one_step( + q=q.squeeze(0), + k=k.squeeze(0), + v=v.squeeze(0), + g=g.squeeze(0), + beta=beta.squeeze(0), + lamb=lamb, + prev_h_kk=last_h_kk, + prev_h_kv=last_h_kv, + max_CG_iteration=self.max_cg_step_decoding, + ) + o = o.unsqueeze(0).to(q) + + update_layer_cache( + self, + past_key_values, + recurrent_state=(h_kk, h_kv), + conv_state=(conv_state_q, conv_state_k), + offset=q_len, + ) + if self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + return o, None, past_key_values diff --git a/fla/layers/mla.py b/fla/layers/mla.py new file mode 100644 index 0000000000000000000000000000000000000000..3cacbb49dca53c1828a2af542dcaacf7ad5b62ad --- /dev/null +++ b/fla/layers/mla.py @@ -0,0 +1,225 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" Implementing the Deepseek Multi Latent Attention (MLA) module. Reference: + +https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py#L328 +""" + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat +from transformers.utils import logging + +from fla.layers.utils import pad_input, unpad_input +from fla.modules import RMSNorm, RotaryEmbedding +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class MultiheadLatentAttention(nn.Module): + r""" + Multi-headed attention from [Deepseek V2](https://arxiv.org/abs/2405.04434) + """ + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 16, + q_lora_rank: int | None = 1536, # q lora rank is optional, None indicates no q lora + qk_rope_head_dim: int = 64, + kv_lora_rank: int = 512, # following the original Deepseek paper + v_head_dim: int = 128, + qk_nope_head_dim: int = 128, + qk_head_dim: int | None = 192, # qk_nope_head_dim + qk_rope_head_dim + window_size: int | None = None, + rope_theta: float = 10000., + max_position_embeddings: int | None = None, + rope_scaling: dict | None = None, + layer_idx: int = None, + ) -> MultiheadLatentAttention: + super().__init__() + + # sanity check + if qk_head_dim is not None: + assert qk_head_dim == qk_nope_head_dim + qk_rope_head_dim, \ + f"qk_head_dim {qk_head_dim} != qk_nope_head_dim {qk_nope_head_dim} + qk_rope_head_dim {qk_rope_head_dim}" + else: + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + # attention params info + self.hidden_size = hidden_size + self.num_heads = num_heads + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.kv_lora_rank = kv_lora_rank + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = qk_head_dim + + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + if q_lora_rank is not None: + self.q_proj = nn.Sequential( + nn.Linear(hidden_size, q_lora_rank, bias=False), + RMSNorm(q_lora_rank, dtype=torch.float32), + nn.Linear(q_lora_rank, self.num_heads * self.qk_head_dim, bias=False), + ) + else: + self.q_proj = nn.Linear(hidden_size, self.num_heads * self.qk_head_dim, bias=False) + + self.k_rope = nn.Linear(hidden_size, self.qk_rope_head_dim, bias=False) + self.kv_proj = nn.Sequential( + nn.Linear(hidden_size, self.kv_lora_rank, bias=False), + RMSNorm(self.kv_lora_rank, dtype=torch.float32), + nn.Linear(self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False), + ) + + self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, hidden_size, bias=False) + + self.scaling = self.qk_head_dim ** (-0.5) + if rope_scaling is not None and rope_scaling.get("rope_type", "default") != "default": + mscale_all_dim = rope_scaling.get("mscale_all_dim", 0) + scaling_factor = rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.scaling = self.scaling * mscale * mscale + + self.rotary = RotaryEmbedding(dim=self.qk_rope_head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + # if attention_mask is not None, this is doing inference + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # prepare q, k, v + batch_size, q_len, _ = hidden_states.shape + + q_states = self.q_proj(hidden_states) + q_states = rearrange(q_states, '... (h d) -> ... h d', d=self.qk_head_dim) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + k_pass, k_rot = self.kv_proj(hidden_states), self.k_rope(hidden_states) + + k_rot = rearrange(k_rot, 'b t d -> b t 1 d') + k_pass = rearrange(k_pass, '... (h d) -> ... h d', d=self.qk_nope_head_dim + self.v_head_dim) + k_pass, v = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + # apply rotary position embedding + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q_len + seqlen_offset + + if attention_mask is not None: + seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1] + max_seqlen = q_len + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + cu_seqlens = kwargs.get("cu_seqlens") + q_rot, k_rot = self.rotary( + q_rot, k_rot, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens, + ) + + k_rot = repeat(k_rot, 'b t 1 d -> b t h d', h=self.num_heads) + q = torch.cat((q_pass, q_rot), dim=-1) + k = torch.cat((k_pass, k_rot), dim=-1) + + # TODO: instead of caching the full k, v, we can actually only cache the compressed_kv and k_rot + # and recover the full k, v from compressed_kv and k_rot + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k, v), + layer_idx=self.layer_idx, + offset=q_len, + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + + # Head dim match to use flash-attn + if self.qk_head_dim != self.v_head_dim: + v = F.pad(v, [0, self.qk_head_dim - self.v_head_dim]) + + # Contains at least one padding token in the sequence + if attention_mask is not None: + if q.shape[1] == 1 and self.window_size is not None: + attention_mask = attention_mask[:, -self.window_size:] + q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len) + cu_seqlens_q, cu_seqlens_k = cu_seqlens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = pad_input(o, indices_q, batch_size, q_len) + elif cu_seqlens is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ).unsqueeze(0) + else: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + + if self.qk_head_dim != self.v_head_dim: + o = o[:, :, :, :self.v_head_dim] + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + return o, None, past_key_values diff --git a/fla/layers/mom.py b/fla/layers/mom.py new file mode 100644 index 0000000000000000000000000000000000000000..06a4e3b2c14aed03c861d7db17f3f44ea9e59548 --- /dev/null +++ b/fla/layers/mom.py @@ -0,0 +1,831 @@ + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, unpad_input, update_layer_cache + + +def _upad_input( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + gate_layer: torch.Tensor, + beta_layer: torch.Tensor, + attention_mask: torch.Tensor, +): + """ + Unpads query, key, and values tensors, using a single dimension for all tokens even though they belong to + different batches. + + This function is used instead of `flash_attn.bert_padding.unpad_input` in order to avoid the recomputation + of the same intermediary + tensors for query, key, value tensors. + + Arguments: + query_layer (`torch.Tensor`): + Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + query_length (`int`): + Target length. + + Return: + query_layer (`torch.Tensor`): + Query state without padding. Shape: (total_target_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + indices_q (`torch.Tensor`): + The indices of non-masked tokens from the flattened input target sequence. + (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index + into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, + `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + query_length = query_layer.shape[1] + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask) + batch_size, kv_seq_len, dim = key_layer.shape + v_dim = value_layer.shape[-1] + + key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, dim), indices_k) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, v_dim), indices_k, + ) + gate_layer = index_first_axis(gate_layer.reshape(batch_size * kv_seq_len, -1), indices_k) + beta_layer = index_first_axis(beta_layer.reshape(batch_size * kv_seq_len, -1), indices_k) + if query_length == kv_seq_len: + query_layer = index_first_axis(query_layer.reshape(batch_size * kv_seq_len, dim), indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device, + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + gate_layer, + beta_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +def transform( + x: torch.Tensor, + routing_mask: torch.Tensor, + num_memories: int, + selected_memories: torch.Tensor, + attention_mask: torch.Tensor, +): + """ + Reorganize token embeddings into memory-aligned chunks. + + Steps: + - Expand for top-k routing if needed. + - Mask out padded tokens via `attention_mask`. + - Sort tokens by (batch, memory). + - Gather and pad tokens per memory slot. + + Args: + x: (batch, seq, hidden) input embeddings. + routing_mask: (batch, seq, num_memories) binary routing mask. + num_memories: number of memory slots. + selected_memories: memory indices per token, + (batch, seq) if k=1 else (batch, seq, topk). + attention_mask: (batch, seq) valid-token mask. + + Returns: + transformed_x: (num_memories, batch, max_len, hidden) reorganized tokens. + truncation_indices: (batch*num_memories, max_len) gather indices. + sorted_indices: (batch*seq*topk,) global sort order. + max_len: int, max tokens per memory. + mask: (batch*num_memories, max_len) validity mask. + mask_2: (num_memories, batch, max_len) validity mask reshaped. + """ + if selected_memories.dim() == 3: + # (batch, seq, topk) + topk = selected_memories.shape[2] + # x (batch, seq, hidden) + x = x.repeat_interleave(topk, dim=1) + # x (batch, seq * topk, hidden) + # (batch, seq, topk) + selected_memories = selected_memories.reshape(selected_memories.shape[0], -1) + # (batch, seq * topk) + + if attention_mask is not None: + attention_mask = attention_mask[:, -routing_mask.shape[1]:] + # mask out the masked tokens + routing_mask[attention_mask.bitwise_not().unsqueeze(-1).expand(-1, -1, num_memories)] = 0 + + b, s, d = x.shape + x_flat = x.reshape(b * s, d) # [b*s, d] + + with torch.no_grad(): + batch_indices = torch.arange(b, device=x.device).unsqueeze(-1) + batch_indices = batch_indices.repeat(1, s).reshape(-1) + if attention_mask is not None: + # sort the masked tokens to the end + batch_indices[attention_mask.repeat_interleave(topk, dim=1).bitwise_not().flatten()] = b + # (b * s) + memories_flat = selected_memories.reshape(-1) # [b*s] + + combined = batch_indices * (memories_flat.max() + 1) + memories_flat + sorted_indices = combined.argsort() + + x_sorted = x_flat[sorted_indices] # [b*s, d] + # (b*s, hidden) -> (b, s, hidd) + with torch.no_grad(): + # routing_mask (b, s, num_memories) + batch_memory_tokens = routing_mask.sum(dim=1) + # (b, num_memories) + flatten_offset = batch_memory_tokens.flatten().cumsum(dim=0) + max_len = batch_memory_tokens.max() + indices = ( + torch.arange(max_len, device=flatten_offset.device).unsqueeze(0).expand(b * num_memories, -1) + + torch.cat([torch.tensor([0], device=flatten_offset.device), flatten_offset[:-1]], dim=0).unsqueeze(1) + ) + mask = indices < flatten_offset.unsqueeze(-1) + truncation_indices = torch.where(mask, indices, torch.zeros_like(indices)) + + gathered_x = torch.gather(x_sorted, 0, truncation_indices.reshape(-1).unsqueeze(-1).expand(-1, d)) + transformed_x = gathered_x.reshape(b * num_memories, -1, d).reshape((b, num_memories, max_len, d)).transpose(0, 1) + # transformed_x = transformed_x * mask.unsqueeze(-1).expand_as(transformed_x) + # pad_x = torch.zeros((b * num_memories, capacity_len-max_len, d), dtype=transformed_x.dtype, device=transformed_x.device) + # pad_mask = torch.zeros((b * num_memories, capacity_len-max_len), dtype=transformed_x.dtype, device=transformed_x.device) + # left pad + # transformed_x = torch.cat((pad_x, transformed_x), dim=1).reshape((b, num_memories, capacity_len, d)).transpose(0, 1) + mask_2 = mask.reshape((b, num_memories, max_len)).transpose(0, 1) + # truncation_indices += capacity_len-max_len + # if attention_mask is not None: + # mask_2 + + return transformed_x, truncation_indices, sorted_indices, max_len, mask, mask_2 + + +def reconstruct( + transformed_x, + indices: torch.Tensor, + sorted_indices: torch.Tensor, + batch_size: int, + seq_len: int, + topk: int, + routing_weights: torch.Tensor, + mask: torch.Tensor, +): + ''' + Reconstruct and mix transformed outputs back into the original input sequence shape. + + Key operations: + 1. Reshapes and transposes `transformed_x` to prepare for scattering. + 2. Applies the `mask` to zero out invalid positions. + 3. Uses `torch.scatter_add_` to scatter and sum the transformed outputs back to their original positions + based on `indices`. + 4. Rearranges the scattered outputs using `sorted_indices` to ensure correct ordering. + 5. Applies the `routing_weights` to weight the outputs. + 6. Sums over the `topk` dimension to produce the final reconstructed output. + + Args: + transformed_x (torch.Tensor): + The transformed output tensor from memory units or experts. + Shape: (num_memories, batch_size, capacity_len, hidden_size) + indices (torch.Tensor): + Indices used for scattering the transformed outputs back to their corresponding positions. + Shape: (batch*num_memories, max_len) + sorted_indices (torch.Tensor): + Sorting indices used to rearrange the scattered outputs back into the original sequence order. + Shape: (batch_size*seq_len*topk) + batch_size (int): + The size of the batch. + seq_len (int): + The length of the input sequence. + topk (int): + The number of top elements selected (`topk`) per token during the selection process. + routing_weights (torch.Tensor): + Routing weights assigned to the top-k selected outputs when reconstructing the final output. + Shape: (batch_size, seq_len, topk) + mask (torch.Tensor): + Boolean mask indicating valid positions in the sequence. + Shape: (batch*num_memories, max_len) + + Returns: + restored_x (torch.Tensor): + The reconstructed output tensor in the original input sequence shape. + Shape: (batch_size, seq_len, hidden_size) + ''' + transformed_x = transformed_x.transpose(0, 1).reshape( + (-1, transformed_x.shape[2], transformed_x.shape[3])) + b, s, k, d = batch_size, seq_len, topk, transformed_x.shape[2] + gathered_x = transformed_x.reshape( + (transformed_x.shape[0] * transformed_x.shape[1], transformed_x.shape[2])) + mask_expanded = mask.reshape(-1).unsqueeze(-1).expand_as(gathered_x) + gathered_x = gathered_x * mask_expanded + + assert (indices >= 0).all(), "Indices should be non-negative" + + resortd_x = torch.zeros((b * s * k, d), device=gathered_x.device, dtype=gathered_x.dtype).scatter_add_( + 0, + indices.reshape(-1).unsqueeze(-1).expand(-1, d), + gathered_x, + ) + assert (indices < resortd_x.size(0)).all(), "Indices should be less than resortd_x size" + + inverse_indices = sorted_indices.argsort() + rearranged_x_flat = resortd_x[inverse_indices] + restored_x = rearranged_x_flat.reshape((b, s * k, d)) + restored_x = restored_x.reshape(b, s, k, d) * routing_weights.reshape(b, s, k).unsqueeze(-1) + restored_x = restored_x.sum(dim=2) + return restored_x + + +class MomAttention(nn.Module): + """ + The layer implementaion for [MoM: Linear Sequence Modeling with Mixture-of-Memories](https://arxiv.org/abs/2502.13685). + """ + + def __init__( + self, + hidden_size: int = 2048, + head_dim: int = 256, + num_heads: int = 4, + expand_v: float = 2, + mode: str = 'chunk', + use_output_gate: bool = True, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + num_memories: int = 8, + topk: int = 2, + capacity: float = 1.0, + shared_mem: bool = False, + single_kv_proj: bool = False, + **kwargs, + ) -> MomAttention: + super().__init__() + self.num_memories = num_memories + self.topk = topk + self.capacity = capacity + self.shared_mem = shared_mem + self.single_kv_proj = single_kv_proj + + self.mode = mode + + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + + self.key_dim = int(self.num_heads * self.head_dim) + self.value_dim = int(self.key_dim * self.expand_v) + self.head_qk_dim = head_dim + self.head_v_dim = int(head_dim * self.expand_v) + self.layer_idx = layer_idx + self.silu = nn.SiLU() + + assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.gate = nn.Linear(self.hidden_size, self.num_memories, bias=False) + if self.single_kv_proj: + self.shared_k = nn.Linear(hidden_size, self.key_dim, bias=False) + self.shared_v = nn.Linear(hidden_size, self.value_dim, bias=False) + self.shared_b = nn.Linear(hidden_size, self.num_heads, bias=False) + self.shared_a = nn.Linear(hidden_size, self.num_heads, bias=False) + else: + self.k_proj = nn.ModuleList([ + nn.Linear(self.hidden_size, self.key_dim, bias=False) + for _ in range(self.num_memories) + ]) + self.v_proj = nn.ModuleList([ + nn.Linear(self.hidden_size, self.value_dim, bias=False) + for _ in range(self.num_memories) + ]) + self.b_proj = nn.ModuleList([ + nn.Linear(self.hidden_size, self.num_heads, bias=False) + for _ in range(self.num_memories) + ]) + self.a_proj = nn.ModuleList([ + nn.Linear(self.hidden_size, self.num_heads, bias=False) + for _ in range(self.num_memories) + ]) + if self.shared_mem: + self.shared_k = nn.Linear(hidden_size, self.key_dim, bias=False) + self.shared_v = nn.Linear(hidden_size, self.value_dim, bias=False) + self.shared_b = nn.Linear(hidden_size, self.num_heads, bias=False) + self.shared_a = nn.Linear(hidden_size, self.num_heads, bias=False) + + A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + raise UserWarning( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + self.apply(self._initialize_weights) + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5) + if module.bias is not None: + nn.init.zeros_(module.bias) + module._is_hf_initialized = True + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + attention_mask = (attention_mask == 1) + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + origin_cu_seqlens = kwargs.get('cu_seqlens') + if origin_cu_seqlens is not None: + hidden_states, attention_mask = self.cu2pad(hidden_states, origin_cu_seqlens) + + mode = 'fused_recurrent' if (hidden_states.shape[1] <= 64 and not self.training) else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = get_layer_cache(self, past_key_values) + # _, q_len = hidden_states.shape[0], hidden_states.shape[1] + + # 🔍 topk gating + router_logits = self.gate(hidden_states) # (bsz, q_len, num_memories) + scores = F.softmax(router_logits, dim=2, dtype=torch.float) + routing_weights, selected_memories = torch.topk(scores, self.topk, dim=-1) # (bsz, seq, topk) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.to(hidden_states.dtype) # we cast back to the input dtype + routing_weights_full = torch.zeros( + routing_weights.shape[0], + routing_weights.shape[1], + self.num_memories, + dtype=routing_weights.dtype, + device=routing_weights.device, + ).scatter(-1, selected_memories, routing_weights) + routing_mask = routing_weights_full.bool().int() + + # if self.use_output_gate: + # o_g = self.g_proj(hidden_states) + + batch_size, seq_len = hidden_states.shape[0], hidden_states.shape[1] + + shared_hidden_states = hidden_states + hidden_states, indices, sorted_indices, max_len, mask, mask_2 = transform( + hidden_states, routing_mask, self.num_memories, selected_memories, attention_mask) + + q = self.q_proj(hidden_states) + if self.single_kv_proj: + k = self.shared_k(hidden_states) + v = self.shared_v(hidden_states) + beta = self.shared_b(hidden_states).sigmoid() + g = -self.A_log.float().exp() * F.softplus(self.shared_a(hidden_states).float() + self.dt_bias) + else: + k = torch.stack([k_expert(hidden_states[i]) for i, k_expert in enumerate(self.k_proj)], dim=0) + v = torch.stack([v_expert(hidden_states[i]) for i, v_expert in enumerate(self.v_proj)], dim=0) + beta = torch.stack([b_expert(hidden_states[i]).sigmoid() for i, b_expert in enumerate(self.b_proj)], dim=0) + g = torch.stack([-self.A_log.float().exp() * F.softplus(a_expert(hidden_states[i]).float() + self.dt_bias) + for i, a_expert in enumerate(self.a_proj)], dim=0) + + q, k, v, g, beta, mask_2 = (rearrange(x, 'e b l ... -> (e b) l ...') for x in (q, k, v, g, beta, mask_2)) + cu_q, cu_k, cu_v, cu_g, cu_beta, indices_q, cu_seqlen_all, max_seq_lens = _upad_input(q, k, v, g, beta, mask_2) + cu_seqlens, reverse_indices = cu_seqlen_all[0].to(torch.long).unique(return_inverse=True) + cu_q, cu_k, cu_v, cu_g, cu_beta = (x.unsqueeze(0).contiguous() for x in (cu_q, cu_k, cu_v, cu_g, cu_beta)) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = [None, None], [None, None], [None, None] + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + + conv_cu_seqlens = cu_seqlens + padded = False + if self.training: + conv_cu_seqlens = None + elif seq_len != 1 and (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() < self.conv_size: + padded = True + conv_cu_seqlens, cu_q, cu_k, cu_v, pad_lengths = self.pad_for_conv(cu_seqlens, cu_q, cu_k, cu_v) + + conv_q = self.prepare_recurrent_state( + conv_state_q[0], + conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + batch_size, + ) + cu_q, conv_q_new = self.q_conv1d( + x=cu_q, + cache=conv_q, + output_final_state=use_cache, + cu_seqlens=conv_cu_seqlens, + ) + conv_state_q[0] = self.handle_recurrent_state( + conv_state_q[0], + conv_q_new, + conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + ) + conv_k = self.prepare_recurrent_state( + conv_state_k[0], + conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + batch_size, + ) + cu_k, conv_k_new = self.k_conv1d( + x=cu_k, + cache=conv_k, + output_final_state=use_cache, + cu_seqlens=conv_cu_seqlens, + ) + conv_state_k[0] = self.handle_recurrent_state( + conv_state_k[0], + conv_k_new, + conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + ) + conv_v = self.prepare_recurrent_state( + conv_state_v[0], + conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + batch_size, + ) + cu_v, conv_v_new = self.v_conv1d( + x=cu_v, + cache=conv_v, + output_final_state=use_cache, + cu_seqlens=conv_cu_seqlens, + ) + conv_state_v[0] = self.handle_recurrent_state( + conv_state_v[0], + conv_v_new, conv_cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + ) + + if padded: + cu_q, cu_k, cu_v = self.unpad_after_conv(conv_cu_seqlens, cu_seqlens, cu_q, cu_k, cu_v, pad_lengths) + + else: + q, k, v = self.silu(q), self.silu(k), self.silu(v) + + cu_q, cu_k, cu_v = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', h=self.num_heads), (cu_q, cu_k, cu_v)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else [ + None for _ in range(1 + self.shared_mem)] + if mode == 'chunk': + o, recurrent_state_ = chunk_gated_delta_rule( + q=cu_q, + k=cu_k, + v=cu_v, + g=cu_g, + beta=cu_beta, + initial_state=recurrent_state[0], + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + recurrent_state[0] = self.handle_recurrent_state( + recurrent_state[0], + recurrent_state_, + cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + ) + + elif mode == 'fused_recurrent': + memories = self.prepare_recurrent_state( + recurrent_state[0], + cu_seqlens, cu_seqlen_all[0], + reverse_indices, batch_size, + ) + o, recurrent_state_ = fused_recurrent_gated_delta_rule( + q=cu_q, + k=cu_k, + v=cu_v, + g=cu_g, + beta=cu_beta, + initial_state=memories, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + recurrent_state[0] = self.handle_recurrent_state( + recurrent_state[0], + recurrent_state_, + cu_seqlens, + cu_seqlen_all[0], + reverse_indices, + ) + + o = o.squeeze(0).contiguous() + o = pad_input(o, indices_q, batch_size*self.num_memories, max_len) + o = rearrange(o, '(e b) l h d -> e b l (h d)', b=batch_size) + o = reconstruct(o, indices=indices, sorted_indices=sorted_indices, batch_size=batch_size, + seq_len=seq_len, topk=self.topk, routing_weights=routing_weights, mask=mask) + o = rearrange(o, 'b l (h d) -> b l h d', h=self.num_heads) + + if self.shared_mem: + shared_o = self.shared_o(shared_hidden_states, attention_mask, recurrent_state, + use_cache, conv_state_q, conv_state_k, conv_state_v) + o += shared_o + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q.shape[2], + ) + + if self.use_output_gate: + g = rearrange(self.g_proj(shared_hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + + if origin_cu_seqlens is not None: + indices, _, _ = get_unpad_data(attention_mask[:, -seq_len:]) + o = index_first_axis(rearrange(o, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + return o, None, past_key_values, router_logits.view(-1, self.num_memories) + + def shared_o( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + recurrent_state=None, + use_cache: bool | None = False, + conv_state_q=[None, None], + conv_state_k=[None, None], + conv_state_v=[None, None], + **kwargs, + ) -> torch.Tensor: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + cu_seqlens = None + if attention_mask is not None: + batch_size, q_len = hidden_states.shape[0], hidden_states.shape[1] + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + q, conv_state_q[1] = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q[1], + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k[1] = self.k_conv1d( + x=self.shared_k(hidden_states), + cache=conv_state_k[1], + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v[1] = self.v_conv1d( + x=self.shared_v(hidden_states), + cache=conv_state_v[1], + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.silu(self.q_proj(hidden_states)) + k = self.silu(self.shared_k(hidden_states)) + v = self.silu(self.shared_v(hidden_states)) + + q, k, v = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', h=self.num_heads), (q, k, v)) + beta = self.shared_b(hidden_states).sigmoid() + g = -self.A_log.float().exp() * F.softplus(self.shared_a(hidden_states).float() + self.dt_bias) + + if mode == 'chunk': + o, recurrent_state[-1] = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state[-1], + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + elif mode == 'fused_recurrent': + o, recurrent_state[-1] = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state[-1], + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + return o + + def cu2pad(self, x, cu_seqlens): + batch_size = cu_seqlens.shape[0] - 1 + max_len = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + indices = torch.tensor([], dtype=torch.long, device=x.device) + attention_mask = torch.ones((batch_size, max_len), dtype=torch.bool, device=x.device) + for i in range(batch_size): + seq_len = cu_seqlens[i+1] - cu_seqlens[i] + pad_len = max_len - seq_len + batch_indices = torch.arange(pad_len, max_len, device=x.device) + batch_indices = batch_indices + i * max_len + indices = torch.cat([indices, batch_indices]) + attention_mask[i, :pad_len] = False + x = pad_input(x.squeeze(0), indices, batch_size, max_len) + return x, attention_mask + + def pad_for_conv(self, cu_seqlens, cu_q, cu_k, cu_v): + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + pad_lengths = torch.clamp(self.conv_size - lengths, min=0) + new_lengths = lengths + pad_lengths + new_cu_seqlens = torch.cat([ + torch.tensor([0], device=cu_seqlens.device, dtype=cu_seqlens.dtype), + torch.cumsum(new_lengths, dim=0), + ]) + final_total_len = new_cu_seqlens[-1].item() + new_q = torch.zeros((1, final_total_len, cu_q.shape[-1]), dtype=cu_q.dtype, device=cu_q.device) + new_k = torch.zeros((1, final_total_len, cu_k.shape[-1]), dtype=cu_k.dtype, device=cu_k.device) + new_v = torch.zeros((1, final_total_len, cu_v.shape[-1]), dtype=cu_v.dtype, device=cu_v.device) + num_sequences = len(lengths) + for i in range(num_sequences): + src_start = cu_seqlens[i] + src_end = cu_seqlens[i+1] + dest_start = new_cu_seqlens[i] + pad_lengths[i] + dest_end = new_cu_seqlens[i+1] + new_q[:, dest_start:dest_end, ...] = cu_q[:, src_start:src_end, ...] + new_k[:, dest_start:dest_end, ...] = cu_k[:, src_start:src_end, ...] + new_v[:, dest_start:dest_end, ...] = cu_v[:, src_start:src_end, ...] + + return new_cu_seqlens, new_q, new_k, new_v, pad_lengths + + def unpad_after_conv(self, conv_cu_seqlens, cu_seqlens, cu_q, cu_k, cu_v, pad_lengths): + original_total_len = cu_seqlens[-1].item() + orig_q = torch.empty((1, original_total_len, cu_q.shape[-1]), dtype=cu_q.dtype, device=cu_q.device) + orig_k = torch.empty((1, original_total_len, cu_k.shape[-1]), dtype=cu_k.dtype, device=cu_k.device) + orig_v = torch.empty((1, original_total_len, cu_v.shape[-1]), dtype=cu_v.dtype, device=cu_v.device) + + num_sequences = len(pad_lengths) + for i in range(num_sequences): + dest_start = cu_seqlens[i] + dest_end = cu_seqlens[i+1] + src_start = conv_cu_seqlens[i] + pad_lengths[i] + src_end = conv_cu_seqlens[i+1] + + orig_q[:, dest_start:dest_end, ...] = cu_q[:, src_start:src_end, ...] + orig_k[:, dest_start:dest_end, ...] = cu_k[:, src_start:src_end, ...] + orig_v[:, dest_start:dest_end, ...] = cu_v[:, src_start:src_end, ...] + return orig_q, orig_k, orig_v + + def prepare_recurrent_state(self, recurrent_state, cu_seqlens, cu_seqlen_all, reverse_indices, batch_size): + if recurrent_state is None: + return None + + if cu_seqlens is None: + return recurrent_state + + total_len = len(cu_seqlen_all) + if len(cu_seqlens) != total_len: + # select memories that are activated + memories = torch.zeros_like(recurrent_state[:self.topk*batch_size]) + mem_id = 0 + for i in range(total_len-1): + if cu_seqlen_all[i] != cu_seqlen_all[i+1]: + memories[mem_id] = recurrent_state[i] + mem_id += 1 + assert mem_id == self.topk * batch_size, f"The number of memories {mem_id} is not correct." + else: + memories = recurrent_state + + return memories + + def handle_recurrent_state(self, recurrent_state, recurrent_state_new, cu_seqlens, cu_seqlen_all, reverse_indices): + if recurrent_state_new is None: + return None + if cu_seqlens is None: + return recurrent_state_new + if recurrent_state is None: + recurrent_state = torch.zeros_like(recurrent_state_new[reverse_indices[1:]-1]) + total_len = len(cu_seqlen_all) + if len(cu_seqlens) != total_len: + for i in range(total_len-1): + if cu_seqlen_all[i] != cu_seqlen_all[i+1]: + recurrent_state[i] = recurrent_state_new[reverse_indices[i+1]-1] + else: + recurrent_state = recurrent_state_new + return recurrent_state diff --git a/fla/layers/multiscale_retention.py b/fla/layers/multiscale_retention.py new file mode 100644 index 0000000000000000000000000000000000000000..cbc8135c7636d4d3cab47ab83a73cc9219a5ca86 --- /dev/null +++ b/fla/layers/multiscale_retention.py @@ -0,0 +1,303 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from transformers.activations import ACT2FN + +from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.rotary import RotaryEmbedding +from fla.ops.retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class MultiScaleRetention(nn.Module): + r""" + The layer implementaion for [Retentive Network: A Successor to Transformer for Large Language Models](https://arxiv.org/pdf/2307.08621.pdf). # noqa + + Args: + mode (str, Optional): + Which Retention kernel to use. + Currently available: `chunk`, `fused_recurrent`, `parallel`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + num_heads (int, Optional): + The number of heads. Default: 8. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1.0, + expand_v: float = 2.0, + num_heads: int = 8, + num_kv_heads: int | None = None, + feature_map: str | None = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + use_output_gate: bool = True, + gate_fn: str = 'swish', + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + fuse_norm: bool = True, + layer_idx: int = None, + **kwargs, + ) -> MultiScaleRetention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.use_output_gate = use_output_gate + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_chunk', 'parallel', 'fused_recurrent'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + if self.use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + if gate_fn == 'swish' and fuse_norm and use_output_gate: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + dtype=torch.float32 + ) + self.gate_fn = ACT2FN[gate_fn] + + # TODO: fix this issue + # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py#L180 + # Ideally, we would want to support arbitrary d_head_qk + assert self.head_k_dim <= 256, "head_k_dim must be less than or equal to 256" + self.rotary = RotaryEmbedding(dim=self.head_k_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if q_len <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim) + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + + seqlen_offset, max_seqlen = 0, q.shape[1] + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None and seqlen_offset > 0: + # to deliminate the offsets of padding tokens + seqlen_offset = prepare_lens_from_mask(attention_mask) - q_len + max_seqlen = q.shape[1] + seqlen_offset.max().item() + + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if self.num_kv_groups > 1: + k = repeat(k, '... h d -> ... (h g) d', g=self.num_kv_groups) + v = repeat(v, '... (h d) -> ... (h g) d', d=self.head_v_dim, g=self.num_kv_groups) + else: + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'parallel': + o, recurrent_state = parallel_retention( + q=q, + k=k, + v=v, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_output_gate: + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, '... (h d) -> ... h d', d=self.head_v_dim) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, '... h d -> ... (h d)') + else: + o = rearrange(self.g_norm(o), '... h d -> ... (h d)') + o = o * self.gate_fn(g) + else: + o = rearrange(self.g_norm(o), '... h d -> ... (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/layers/nsa.py b/fla/layers/nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..d21bea41083950dd1ceb294bd15d4e46b4901dc0 --- /dev/null +++ b/fla/layers/nsa.py @@ -0,0 +1,137 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.utils import logging + +from fla.modules import RotaryEmbedding +from fla.ops.nsa.parallel import parallel_nsa +from fla.ops.utils.index import prepare_lens_from_mask + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class NativeSparseAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 64, + num_kv_heads: int | None = 4, + head_dim: int = 64, + qkv_bias: bool = False, + block_size: int | None = 64, + block_counts: torch.LongTensor | int | None = 16, + window_size: int | None = 512, + rope_theta: float | None = 10000., + max_position_embeddings: int | None = None, + layer_idx: int = None, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + + self.block_size = block_size + self.block_counts = block_counts + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.g_proj = nn.Linear(self.hidden_size, self.num_heads * 3, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, seq_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=3) + g_cmp, g_slc, g_swa = g.sigmoid().unbind(-1) + + cu_seqlens = kwargs.get('cu_seqlens') + + seqlen_offset, max_seqlen = 0, seq_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=seq_len, + cache_kwargs=dict(window_size=self.window_size), + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + o = parallel_nsa( + q=q, + k=k, + v=v, + g_cmp=g_cmp, + g_slc=g_slc, + g_swa=g_swa, + block_size=self.block_size, + block_counts=self.block_counts, + window_size=self.window_size, + cu_seqlens=cu_seqlens, + ) + o = o.reshape(batch_size, seq_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/path_attn.py b/fla/layers/path_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..cc7ddc5abc03c1de7eea226a72205cc68b3cd663 --- /dev/null +++ b/fla/layers/path_attn.py @@ -0,0 +1,216 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.utils import logging + +from fla.layers.utils import pad_input, unpad_input +from fla.modules import RMSNorm, ShortConvolution +from fla.modules.l2norm import l2_norm +from fla.ops.attn.decoding import attn_decoding_one_step +from fla.ops.path_attn.parallel import parallel_path_attn + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class PaTHAttention(nn.Module): + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: int | None = None, + use_forget_gate: bool = False, + use_qk_norm: bool = False, + layer_idx: int = None, + use_low_rank_w: bool = True, + use_w_shortconv: bool = True, + conv_size: int = 3, + conv_bias: bool = False, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False) + + # We use low-rank parameterization for the w_proj to reduce parameters in MHA settings. + if use_low_rank_w: + self.w_proj = nn.Sequential( + nn.Linear(self.hidden_size, 32, bias=False), + nn.Linear(32, self.kv_dim, bias=False), + ) + # In MQA/GQA settings, key/value heads are shared, so we use a standard linear projection + # which doesn't introduce too many parameters + else: + self.w_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False) + + # per head norm + if use_qk_norm: + self.maybe_q_norm = RMSNorm(self.head_dim, dtype=torch.float32) + self.maybe_k_norm = RMSNorm(self.head_dim, dtype=torch.float32) + else: + self.maybe_q_norm = nn.Identity() + self.maybe_k_norm = nn.Identity() + + if use_w_shortconv: + self.w_conv1d = ShortConvolution(hidden_size=self.kv_dim, kernel_size=conv_size, bias=conv_bias, activation='silu') + self.use_w_shortconv = use_w_shortconv + self.bt_proj = nn.Linear(self.hidden_size, self.num_kv_heads, bias=True) + self.use_forget_gate = use_forget_gate + if use_forget_gate: + self.g_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + if use_cache: + assert past_key_values is not None, "past_key_values must be provided when use_cache is True" + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + batch_size, q_len, _ = hidden_states.size() + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + w = self.w_proj(hidden_states) + beta = self.bt_proj(hidden_states).float().sigmoid() * 2 # allowing negative eigenvalues + g = F.logsigmoid(self.g_proj(hidden_states).float()) if self.use_forget_gate else None + cu_seqlens = kwargs.get('cu_seqlens') + assert not (cu_seqlens is not None and attention_mask is not None), ( + "cu_seqlens should not be provided when attention_mask is not None" + ) + # Training + if attention_mask is None: + assert use_cache is False, "use_cache should be False in training" + if self.use_w_shortconv: + w, _ = self.w_conv1d(w, cache=None, output_final_state=False, cu_seqlens=cu_seqlens) + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim) + q, k = self.maybe_q_norm(q), self.maybe_k_norm(k) + w = l2_norm(w, output_dtype=torch.float32) + o, _ = parallel_path_attn(q=q, k=k, v=v, w=w, beta=beta, g=g, cu_seqlens=cu_seqlens) + + # Prefilling or decoding + else: + assert self.training is False, "attention mask is not supported in training. Please use variable length input." + try: + last_state = past_key_values[self.layer_idx] + except KeyError: + last_state = None + # Decoding + if last_state is not None: + if g is not None: + past_k, past_v, past_g = last_state['attn_state'] + else: + past_k, past_v = last_state['attn_state'] + past_g = None + w_conv_state = last_state['conv_state'] + past_k = rearrange(past_k, '... (h d) -> ... h d', d=self.head_dim) + if self.use_w_shortconv: + w, w_conv_state = self.w_conv1d(w, cache=w_conv_state, output_final_state=use_cache, cu_seqlens=cu_seqlens) + w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim) + w = l2_norm(w, output_dtype=torch.float32) + + @torch.compile + def rank_one_update(k, w, beta): + original_dtype = k.dtype + k = k.float() + w = w.float() + beta = beta.float() + k = k - beta[..., None].float() * (k * w).sum(-1, keepdim=True) * w + return k.to(original_dtype) + + past_k = rank_one_update(past_k, w, beta) + past_k = rearrange(past_k, '... h d -> ... (h d)') + k = torch.cat([past_k, k], dim=1) + v = torch.cat([past_v, v], dim=1) + g = torch.cat([past_g, g], dim=1) if g is not None else None + past_key_values[self.layer_idx]['attn_state'] = (k, v, g) if g is not None else (k, v) + past_key_values.update( + conv_state=w_conv_state, + layer_idx=self.layer_idx, + offset=q_len, + ) + if g is not None: + q, (k, v, g), indices_q, cu_seqlens, max_seq_lens = unpad_input( + q, (k, v, g), attention_mask, q_len, keepdim=True) + max_seqlen_q, max_seqlen_k = max_seq_lens + else: + q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input( + q, (k, v), attention_mask, q_len, keepdim=True) + max_seqlen_q, max_seqlen_k = max_seq_lens + _, cu_seqlens = cu_seqlens + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + assert max_seqlen_q == 1, "only support q_len == 1 for decoding" + o = attn_decoding_one_step(q, k, v, g, cu_seqlens=cu_seqlens, do_gate_scale=True) # reduced to fox's decoding + # Prefilling + else: + v_cache = v.clone() + g_cache = g.clone() if g is not None else None + if g is None: + q, (k, v, w, beta), indices_q, cu_seqlens, max_seq_lens = unpad_input( + q, (k, v, w, beta), attention_mask, q_len, keepdim=True) + else: + q, (k, v, w, beta, g), indices_q, cu_seqlens, max_seq_lens = unpad_input( + q, (k, v, w, beta, g), attention_mask, q_len, keepdim=True) + max_seqlen_q, max_seqlen_k = max_seq_lens + assert max_seqlen_q == max_seqlen_k, "max_seqlen_q should be equal to max_seqlen_k in prefilling" + _, cu_seqlens = cu_seqlens + if self.use_w_shortconv: + w, w_conv_state = self.w_conv1d(w, cache=None, output_final_state=use_cache, cu_seqlens=cu_seqlens) + else: + w_conv_state = None + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim) + w = l2_norm(w, output_dtype=torch.float32) + o, k_cache = parallel_path_attn(q=q, k=k, v=v, w=w, beta=beta, g=g, + cu_seqlens=cu_seqlens, use_cache=use_cache) + if use_cache: + k_cache = pad_input(k_cache.squeeze(0), indices_q, batch_size, q_len) + k_cache = rearrange(k_cache, '... h d -> ... (h d)') + past_key_values.update( + attn_state=(k_cache, v_cache, g_cache) if g_cache is not None else (k_cache, v_cache), + conv_state=w_conv_state, + layer_idx=self.layer_idx, + offset=q_len, + ) + o = pad_input(o.squeeze(0), indices_q, batch_size, q_len) + o = rearrange(o, '... h d -> ... (h d)') + o = self.o_proj(o) + return o, None, past_key_values diff --git a/fla/layers/quasar.py b/fla/layers/quasar.py new file mode 100644 index 0000000000000000000000000000000000000000..dc2ae2aff1b906994542d0edebf3c8dc79e53dad --- /dev/null +++ b/fla/layers/quasar.py @@ -0,0 +1,439 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention + +from __future__ import annotations + +import contextlib +import math +import os +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_unpad_data, index_first_axis, pad_input + + +def _quasar_debug_tensor(name: str, tensor: torch.Tensor, layer_idx: int | None) -> None: + if os.environ.get("QUASAR_DEBUG_FINITE", "0") != "1": + return + if tensor is None or torch.isfinite(tensor).all(): + return + with torch.no_grad(): + t = torch.nan_to_num(tensor.detach().float(), nan=0.0, posinf=0.0, neginf=0.0) + nonfinite = int((~torch.isfinite(tensor)).sum().item()) + print( + f"[QUASAR DEBUG] layer={layer_idx} stage={name} nonfinite={nonfinite} " + f"min={float(t.min())} max={float(t.max())} mean={float(t.mean())}", + flush=True, + ) + + +class _TorchRMSNormGated(nn.Module): + def __init__(self, hidden_size: int, activation: str = "sigmoid", eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.activation = activation + self.eps = eps + + def reset_parameters(self) -> None: + self.weight.data.fill_(1.0) + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + y = torch.nan_to_num( + x.float(), + nan=0.0, + posinf=1e4, + neginf=-1e4, + ).clamp_(min=-1e4, max=1e4) + y = y * torch.rsqrt(y.square().mean(dim=-1, keepdim=True) + self.eps) + weight = torch.nan_to_num( + self.weight.float(), + nan=1.0, + posinf=1.0, + neginf=1.0, + ).clamp_(min=0.0, max=4.0) + y = y.to(dtype) * weight.to(dtype=dtype, device=x.device) + gate = torch.nan_to_num( + gate.float(), + nan=0.0, + posinf=30.0, + neginf=-30.0, + ).clamp_(min=-30.0, max=30.0) + if self.activation in {"swish", "silu"}: + gate = gate * torch.sigmoid(gate) + elif self.activation == "sigmoid": + gate = torch.sigmoid(gate) + return y * gate.to(dtype=dtype, device=x.device) + + +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): + """Applies Rotary Position Embedding to the query and key tensors.""" + # cos, sin: [1, 1, seq_len, rotary_dim] + # q, k: [batch_size, seq_len, n_heads, head_dim] + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + cos = cos.transpose(1, 2) # [1, seq_len, 1, rotary_dim] + sin = sin.transpose(1, 2) # [1, seq_len, 1, rotary_dim] + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + return torch.cat([q_embed, q_pass], dim=-1), torch.cat([k_embed, k_pass], dim=-1) + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class QuasarAttention(nn.Module): + """ + QuasarAttention layer implementation. + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + head_dim (int, Optional): + The dimension of each head. Default: 128. + num_heads (int, Optional): + The number of heads. Default: 16. + mode (str, Optional): + Which QuasarAttention kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + head_dim: int = 128, + num_heads: int = 16, + mode: str = "chunk", + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> QuasarAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.key_dim = int(self.num_heads * self.head_dim) + self.value_dim = int(self.num_heads * self.head_dim) + self.layer_idx = layer_idx + + assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + # KDA matching: Use SiLU on q, k, v for better learning if not using short conv + # (Short conv already has its own activation) + self.q_act = nn.SiLU() + self.k_act = nn.SiLU() + self.v_act = nn.SiLU() + + if use_short_conv: + from fla.modules.convolution import ShortConvolution + + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + + # Data-dependent Beta (Adaptive Decay) + # Instead of a static per-head parameter, we use a linear projection + # to allow the model to learn contextual importance (read/write sharpness). + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + + # Learnable state decay (like KDA/Mamba A matrix) + self.A_log = nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16))) + self.A_log._no_weight_decay = True + self.dt_bias = nn.Parameter(torch.zeros(self.key_dim, dtype=torch.float32)) + self.dt_bias._no_weight_decay = True + + # KIMI matches: separate f_proj for kernel and g_proj for final output gating + self.f_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.g_proj = nn.Sequential( + nn.Linear(hidden_size, self.head_dim, bias=False), + nn.Linear(self.head_dim, self.value_dim, bias=True), + ) + + self.o_norm = _TorchRMSNormGated(self.head_dim, activation="sigmoid", eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def reset_parameters(self) -> None: + for module in self.children(): + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + mode = self.mode + if self.training and mode == "fused_recurrent": + # The fused recurrent Quasar path is forward-only in this tree. + # Training must use the chunk kernel until its backward exists. + mode = "chunk" + + # Bailing hidden states can be very large after MoE/FSDP checkpoint + # restore. Quasar's delta-rule triangular solve is much more sensitive + # to projection scale than GQA/GLA, so sanitize and RMS-normalize only + # the Quasar branch input. The residual model path remains untouched. + input_dtype = hidden_states.dtype + hidden_states = torch.nan_to_num( + hidden_states.float(), + nan=0.0, + posinf=60.0, + neginf=-60.0, + ).clamp_(min=-60.0, max=60.0) + hidden_states = hidden_states * torch.rsqrt( + hidden_states.square().mean(dim=-1, keepdim=True) + 1e-6 + ) + hidden_states = hidden_states.to(dtype=input_dtype) + _quasar_debug_tensor("input_normed", hidden_states, self.layer_idx) + + last_state = None + recurrent_state = None + conv_state_q, conv_state_k, conv_state_v = None, None, None + + if past_key_values is not None and self.layer_idx is not None: + if hasattr(past_key_values, "recurrent_states") and self.layer_idx in past_key_values.recurrent_states: + recurrent_state = past_key_values.recurrent_states[self.layer_idx] + if hasattr(past_key_values, "conv_states") and self.layer_idx in past_key_values.conv_states: + conv_state_q, conv_state_k, conv_state_v = past_key_values.conv_states[self.layer_idx] + else: + try: + # Standard list/tuple cache (FLA style fallback) + if len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + if isinstance(last_state, dict): + recurrent_state = last_state.get("recurrent_state", None) + convs = last_state.get("conv_state", None) + if convs is not None: + conv_state_q, conv_state_k, conv_state_v = convs + except TypeError: + pass + + cu_seqlens = kwargs.get("cu_seqlens") + if attention_mask is not None: + # Optimization: Skip unpadding if all tokens are valid (common in packed distillation) + if attention_mask.all(): + indices, cu_seqlens = None, None + else: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + else: + indices = None + + if self.use_short_conv: + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_act(self.q_proj(hidden_states)) + k = self.k_act(self.k_proj(hidden_states)) + v = self.v_act(self.v_proj(hidden_states)) + _quasar_debug_tensor("q_proj", q, self.layer_idx) + _quasar_debug_tensor("k_proj", k, self.layer_idx) + _quasar_debug_tensor("v_proj", v, self.layer_idx) + + q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + # Apply RoPE if provided + cos = kwargs.get("cos") + sin = kwargs.get("sin") + if cos is not None and sin is not None: + if attention_mask is not None: + # Unpad cos/sin using the same indices + # cos/sin shape is [1, 1, seq_len, head_dim] or [batch_size, seq_len, head_dim] + if cos.shape[0] == 1 and cos.shape[1] == 1: + # Broadcastable/Shared RoPE [1, 1, seq_len, head_dim] + # We need to expand to [batch_size, seq_len, head_dim] before unpadding + cos_expanded = cos.squeeze(1).expand(batch_size, -1, -1) + sin_expanded = sin.squeeze(1).expand(batch_size, -1, -1) + cos = index_first_axis(rearrange(cos_expanded, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1) + sin = index_first_axis(rearrange(sin_expanded, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1) + else: + # Already [batch_size, 1, seq_len, head_dim] or [batch_size, seq_len, head_dim] + if cos.dim() == 4: + cos = cos.squeeze(1) + sin = sin.squeeze(1) + cos = index_first_axis(rearrange(cos, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1) + sin = index_first_axis(rearrange(sin, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1) + + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + # QK Normalization AFTER RoPE — ensures kernel receives unit-norm vectors + # regardless of any precision drift introduced by the rotation + q = F.normalize(q, p=2, dim=-1) + k = F.normalize(k, p=2, dim=-1) + _quasar_debug_tensor("q_norm", q, self.layer_idx) + _quasar_debug_tensor("k_norm", k, self.layer_idx) + + # Adaptive Beta: Sigmoid(b_proj(x)) is bounded to (0, 1) to prevent explosions. + beta = self.b_proj(hidden_states).sigmoid() + _quasar_debug_tensor("beta", beta, self.layer_idx) + + if mode == "chunk": + from fla.ops.quasar.chunk import chunk_quasar + + o, recurrent_state = chunk_quasar( + q=q, + k=k, + v=v, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + _quasar_debug_tensor("chunk_kernel_o", o, self.layer_idx) + elif mode == "fused_recurrent": + from fla.ops.quasar.fused_recurrent import fused_recurrent_quasar + + # Use f_proj for kernel gate in fused mode + f_gate = self.f_proj(hidden_states) + f_gate = rearrange(f_gate, "... (h d) -> ... h d", d=self.head_dim) + o, recurrent_state = fused_recurrent_quasar( + q=q, + k=k, + v=v, + g=f_gate, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + ) + _quasar_debug_tensor("fused_kernel_o", o, self.layer_idx) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + o = torch.nan_to_num( + o.float(), + nan=0.0, + posinf=1e4, + neginf=-1e4, + ).clamp_(min=-1e4, max=1e4).to(dtype=v.dtype) + _quasar_debug_tensor("kernel_o_clamped", o, self.layer_idx) + + if past_key_values is not None: + if hasattr(past_key_values, "update_quasar_state"): + past_key_values.update_quasar_state( + self.layer_idx, + recurrent_state, + (conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None + ) + else: + with contextlib.suppress(TypeError): + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q_len, + ) + + # Final output gating using g_proj + # Handle flattened inputs (unpadded) from FSDP/Flash-Linear-Attention + if hidden_states.dim() == 2: + # (N, D) -> (N, H, D/H) + g = self.g_proj(hidden_states) + g = rearrange(g, "n (h d) -> n h d", d=self.head_dim) + _quasar_debug_tensor("output_gate", g, self.layer_idx) + o = self.o_norm(o, g) + o = rearrange(o, "n h d -> n (h d)") + else: + # (B, S, D) -> (B, S, H, D/H) + g = self.g_proj(hidden_states) + g = rearrange(g, "b s (h d) -> b s h d", d=self.head_dim) + _quasar_debug_tensor("output_gate", g, self.layer_idx) + o = self.o_norm(o, g) + o = rearrange(o, "b s h d -> b s (h d)") + _quasar_debug_tensor("post_norm_gate", o, self.layer_idx) + + o = self.o_proj(o) + _quasar_debug_tensor("o_proj", o, self.layer_idx) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + # LFM2 expects 2 return values (hidden_states, _) + return o, None diff --git a/fla/layers/rebased.py b/fla/layers/rebased.py new file mode 100644 index 0000000000000000000000000000000000000000..eaf7ec9dd6e68d96b5c906263f8ab076eb5b9928 --- /dev/null +++ b/fla/layers/rebased.py @@ -0,0 +1,144 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/layers/rebased_fast.py +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.modules.feature_map import RebasedFeatureMap +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn +from fla.ops.rebased import parallel_rebased + + +class ReBasedLinearAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + l_max: int = 2048, + feature_dim: int = 16, + num_key_value_heads: int = 16, + num_heads: int = 16, + use_gamma: bool | None = True, + use_beta: bool | None = True, + normalize: bool | None = True, + causal: bool = True, + eps: float = 1e-5, + mode: str = "parallel", + layer_idx: int | None = None, + **kwargs, + ) -> ReBasedLinearAttention: + super().__init__() + self.hidden_size = hidden_size + self.l_max = l_max + self.mode = mode + assert self.mode in ["fused_chunk", "parallel", 'chunk'] + + self.feature_dim = feature_dim + self.num_key_value_heads = num_key_value_heads + self.num_heads = num_heads + self.head_dim = self.hidden_size // self.num_key_value_heads + self.use_gamma = use_gamma + self.use_beta = use_beta + self.normalize = normalize + self.causal = causal + self.eps = eps + self.mode = mode + self.layer_idx = layer_idx + + self.feature_map = RebasedFeatureMap(self.feature_dim, use_gamma, use_beta, normalize) + self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.dropout = nn.Identity() + + def forward(self, hidden_states: torch.Tensor, **kwargs): + mode = self.mode + q = rearrange( + self.q_proj(hidden_states), + "... (h d) -> ... h d", + h=self.num_heads, + d=self.feature_dim, + ) + k = rearrange( + self.k_proj(hidden_states), + "... (h d) -> ... h d", + h=self.num_heads, + d=self.feature_dim, + ) + v = rearrange( + self.v_proj(hidden_states), + "... (h d) -> ... h d", + h=self.num_key_value_heads, + d=self.head_dim, + ) + q, k = self.feature_map(q, flatten=(mode != 'parallel')), self.feature_map(k, flatten=(mode != 'parallel')) + if mode == "fused_chunk": + o = fused_chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=True, + scale=1, + ) + elif mode == 'chunk': + o = chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=True, + scale=1, + ) + elif mode == 'parallel': + assert q.shape[-1] <= 128 + o = parallel_rebased( + q=q, + k=k, + v=v, + eps=self.eps, + use_scale=True, + use_normalize=True, + ) + o = rearrange(o, "... h d -> ... (h d)") + o = self.o_proj(o) + o = self.dropout(o) + return o + + # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119 + def forward_reference( + self, + hidden_states: torch.Tensor, + filters: torch.Tensor = None, + *args, + **kwargs, + ): + """ + x (torch.Tensor): tensor of shape (b, d, t) + y (torch.Tensor): tensor of shape (b, d, t) + """ + b, t, _ = hidden_states.size() + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + + q = q.view(b, t, -1, self.feature_dim).transpose(1, 2) + k = k.view(b, t, -1, self.feature_dim).transpose(1, 2) + v = v.view(b, t, -1, self.head_dim).transpose(1, 2) + + # Linear attention + q, k = self.feature_map(q), self.feature_map(k) + q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1) + + # Compute attention + if self.causal: + y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps)) + else: + y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps)) + y = rearrange(y, 'b h t d -> b t (h d)') + y = self.o_proj(y.to(hidden_states.dtype)) + y = self.dropout(y) + return y.to(hidden_states.dtype) diff --git a/fla/layers/rodimus.py b/fla/layers/rodimus.py new file mode 100644 index 0000000000000000000000000000000000000000..b016ef8ea7e2c227979e5a81e559ece0805583e4 --- /dev/null +++ b/fla/layers/rodimus.py @@ -0,0 +1,397 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange, repeat +from transformers.utils import logging + +from fla.layers.utils import ( + get_layer_cache, + get_unpad_data, + index_first_axis, + pad_input, + require_cache_layer_idx, + unpad_input, + update_layer_cache, +) +from fla.modules import RMSNorm, RotaryEmbedding, ShortConvolution +from fla.modules.layernorm_gated import RMSNormGated +from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +def align_multiple(value, multiple_size=8): + if value % multiple_size != 0: + value += multiple_size - (value % multiple_size) + return value + + +def autocast_to_fp16(x): + if x.dtype not in {torch.float16, torch.bfloat16}: + return x.to(dtype=torch.bfloat16) + else: + return x + + +class RodimusAttention(nn.Module): + def __init__( + self, + block_type: str = 'rodimus', + mode: str = 'chunk', + hidden_size: int = 1024, + input_gate_low_rank: float | str | None = 'auto', + expand_ratio: int = 64, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = True, + norm_eps: float = 1e-5, + k_norm_eps: float | None = None, + residual_in_fp32: bool = True, + layer_idx: int = None, + ): + super().__init__() + + self.block_type = block_type + self.mode = mode + self.hidden_size = hidden_size + self.d_inner = align_multiple(int(self.hidden_size * 2), 8) + + self.expand_ratio = expand_ratio + self.input_gate_low_rank = max(self.hidden_size // 64, 16) if input_gate_low_rank == "auto" else input_gate_low_rank + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.norm_eps = norm_eps + self.k_norm_eps = k_norm_eps if k_norm_eps is not None else 1e-12 + self.mem_size = expand_ratio + + self.residual_in_fp32 = residual_in_fp32 + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`." + + self.gate_proj = nn.Linear(self.hidden_size, self.d_inner, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.d_inner, bias=False) + self.activation_norm = RMSNormGated(hidden_size=self.d_inner, eps=norm_eps, norm_before_gate=False) + self.down_proj = nn.Linear(self.d_inner, self.hidden_size, bias=False) + + if use_short_conv: + self.short_conv = ShortConvolution( + hidden_size=self.d_inner, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + self.residual_weight = nn.Parameter(torch.ones( + (self.d_inner, ), dtype=torch.float32 if self.residual_in_fp32 else None), requires_grad=True) + + self.k_proj = nn.Linear(self.d_inner, self.mem_size, bias=False) + self.q_proj = nn.Linear(self.d_inner, self.mem_size, bias=False) + + self.g_gate_proj = nn.Linear(self.d_inner, self.mem_size, bias=True) + self.tau_gate_proj = nn.Linear(self.d_inner, self.mem_size, bias=True) + self.i_gate_proj = nn.Sequential( + nn.Linear(self.d_inner, self.input_gate_low_rank, bias=False), + nn.Linear(self.input_gate_low_rank, self.d_inner, bias=True), + nn.Sigmoid(), + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + hidden_states, final_gate = self.up_proj(hidden_states), self.gate_proj(hidden_states) + + if self.use_short_conv: + conv_state = None + if last_state is not None: + conv_state = last_state['conv_state'] + shift_hidden_states, conv_state = self.short_conv( + x=hidden_states, + cache=conv_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + shift_hidden_states = hidden_states + + q = self.q_proj(shift_hidden_states) + k = self.k_proj(shift_hidden_states) + v = self.i_gate_proj(hidden_states) * hidden_states + + g_gate = F.linear(shift_hidden_states, self.g_gate_proj.weight) + self.g_gate_proj.bias.float() + tau_gate = F.linear(shift_hidden_states, self.tau_gate_proj.weight) + self.tau_gate_proj.bias.float() + + g_gate = F.softplus(g_gate) + it_gate = g_gate + rt_gate_log = -g_gate + + tau_gate = F.sigmoid(tau_gate) + it_gate = it_gate ** tau_gate + rt_gate_log = rt_gate_log * tau_gate + + k = F.normalize(k.float(), dim=-1, eps=self.k_norm_eps) * it_gate + q, k, v, rt_gate_log = map(lambda x: x.unsqueeze(1).transpose(1, 2), (q, k, v, rt_gate_log)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=rt_gate_log, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_gla( + q=q, + k=k, + v=v, + g=rt_gate_log, + initial_state=recurrent_state, + output_final_state=use_cache, + head_first=False, + ) + elif mode == 'chunk': + q, k, rt_gate_log = map(lambda x: x.to(v.dtype), (q, k, rt_gate_log)) + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=rt_gate_log, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + rodimus_caches = None + if past_key_values is not None: + if self.block_type == 'rodimus': + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=conv_state if self.use_short_conv else None, + offset=q_len, + ) + else: + rodimus_caches = (recurrent_state, conv_state if self.use_short_conv else None) + + o = (o.transpose(1, 2).squeeze(1) + (shift_hidden_states.float() + if self.residual_in_fp32 else shift_hidden_states) * self.residual_weight).to(o.dtype) + + o = self.activation_norm(o, final_gate) + o = self.down_proj(o) + + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + if self.block_type == 'rodimus': + return o, None, past_key_values + else: + return o, None, (past_key_values, rodimus_caches) + + +class SlidingWindowSharedKeyAttention(nn.Module): + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: int = 2048, + rope_theta: float | None = 10000., + max_position_embeddings: int | None = None, + layer_idx: int = None, + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = self.hidden_size // self.num_heads + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.head_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32) + self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + rodimus_caches = kwargs.get('rodimus_caches') + + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + # equivalent to cu_seqlens in `flash_attn` + cu_seqlens = kwargs.get('cu_seqlens') + + layer_idx = require_cache_layer_idx(self, past_key_values) + seqlen_offset, max_seqlen = 0, q.shape[1] + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + if rodimus_caches is not None: + recurrent_state, conv_state = rodimus_caches + else: + recurrent_state, conv_state = None, None + + cache_has_content = past_key_values.get_seq_length(layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + recurrent_state=recurrent_state, + conv_state=conv_state, + attn_state=[k.flatten(-2, -1), v.flatten(-2, -1)], + layer_idx=layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size), + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + q, k, v = map(autocast_to_fp16, (q, k, v)) + k = repeat(k, "... h d -> ... (n h) d", n=self.num_heads) + # Contains at least one padding token in the sequence + if attention_mask is not None: + q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input( + q=q, + states=(k, v), + attention_mask=attention_mask[:, -max(self.window_size, q_len):], + q_len=q_len, + ) + cu_seqlens_q, cu_seqlens_k = cu_seqlens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = pad_input(o, indices_q, batch_size, q_len) + elif cu_seqlens is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ).unsqueeze(0) + else: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0), + ) + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o.to(dtype=self.o_proj.weight.dtype)) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/rwkv6.py b/fla/layers/rwkv6.py new file mode 100644 index 0000000000000000000000000000000000000000..ac154f30a217ea45ef3deb36039aa877a030ea58 --- /dev/null +++ b/fla/layers/rwkv6.py @@ -0,0 +1,360 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# "Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence"[https://arxiv.org/abs/2404.05892] + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import GroupNorm +from fla.modules.activations import ACT2FN +from fla.modules.token_shift import token_shift +from fla.ops.rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class RWKV6Attention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 0.5, + expand_v: float = 1.0, + num_heads: int = 4, + gate_fn: str = 'swish', + proj_low_rank_dim: int = 32, + gate_low_rank_dim: int = 64, + fuse_norm: bool = True, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + **kwargs, + ) -> RWKV6Attention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.proj_low_rank_dim = proj_low_rank_dim + self.gate_low_rank_dim = gate_low_rank_dim + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + self.x_proj = nn.Sequential( + LerpLinear(hidden_size, proj_low_rank_dim * 5), + nn.Tanh(), + nn.Linear(proj_low_rank_dim * 5, hidden_size, bias=False), + ) + self.x_bias = nn.Parameter(torch.zeros(5, hidden_size)) + + self.r_proj = DDLerpLinear(hidden_size, self.key_dim) + self.w_proj = DDLerpLinear(hidden_size, self.key_dim, low_rank_dim=gate_low_rank_dim) + self.k_proj = DDLerpLinear(hidden_size, self.key_dim) + self.v_proj = DDLerpLinear(hidden_size, self.value_dim) + self.g_proj = DDLerpLinear(hidden_size, self.value_dim) + self.bonus = nn.Parameter(torch.zeros(num_heads, self.head_k_dim)) + + # TODO: fuse GroupNorm and output gate + self.g_norm = GroupNorm(self.num_heads, self.value_dim, elementwise_affine=elementwise_affine, bias=True, eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + self.gate_fn = ACT2FN[gate_fn] + + try: + from transformers.modeling_utils import _init_weights + except ImportError: + _init_weights = True + if _init_weights: + self.apply(self._initialize_weights) + + warnings.warn( + "According to Bo, you are using a potentially buggy FLA implementation of RWKV. " + "If you plan to report any numbers based on this implementation, we strongly recommend " + "cross-checking with the official repo: https://github.com/BlinkDL/RWKV-LM. " + "Bo may disagree with results reported from this version.", + ) + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5) + if module.bias is not None: + nn.init.zeros_(module.bias) + if isinstance(module, nn.Parameter): + nn.init.xavier_uniform_(module, gain=2 ** -2.5) + module._is_hf_initialized = True + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, seq_len, hidden_size = hidden_states.shape + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + if attention_mask is not None: + hidden_states = hidden_states.mul(attention_mask[:, -hidden_states.shape[-2]:, None]) + + if hidden_states.shape[1] == 1 and last_state is not None: + shifted = last_state['conv_state'].unsqueeze(1) + delta = shifted - hidden_states + elif last_state is None: + delta = token_shift(hidden_states, cu_seqlens) + else: + shifted = self.time_shift(hidden_states) + shifted[:, 0] = last_state['conv_state'] + delta = shifted - hidden_states + + x = self.x_proj[0](hidden_states, delta, cu_seqlens).view(batch_size, seq_len, -1, self.proj_low_rank_dim) + x = torch.einsum('b t n r, h n r-> b t n h', self.x_proj[1](x), self.x_proj[2].weight.view(hidden_size, 5, -1)) + + r, w, k, v, g = x.add_(self.x_bias).unbind(-2) + r = self.r_proj(hidden_states, r, delta, cu_seqlens) + w = self.w_proj(hidden_states, w, delta, cu_seqlens) + k = self.k_proj(hidden_states, k, delta, cu_seqlens) + v = self.v_proj(hidden_states, v, delta, cu_seqlens) + g = self.g_proj(hidden_states, g, delta, cu_seqlens) + + # dealing with left-padding + if attention_mask is not None: + v = v.mul(attention_mask[:, -v.shape[-2]:, None]) + r, w, k = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_k_dim), (r, w, k)) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + w = -torch.exp(w) + u = self.bonus + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_rwkv6( + r=r, + k=k, + v=v, + w=w, + u=u, + scale=1., + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_rwkv6( + r=r, + k=k, + v=v, + w=w, + u=u, + scale=1., + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=hidden_states[:, -1], + offset=seq_len, + ) + + o = self.g_norm(rearrange(o, '... h d -> ... (h d)')) * self.gate_fn(g) + o = self.o_proj(o) + + return o, None, past_key_values + + +class LoRA(nn.Module): + + def __init__( + self, + input_dim: int, + output_dim: int, + low_rank_dim: int, + bias: bool | None = True, + activation: str | None = 'tanh', + ): + super().__init__() + + self.input_dim = input_dim + self.output_dim = output_dim + self.low_rank_dim = low_rank_dim + self.bias = bias + + if activation is None: + self.activation = nn.Identity() + elif activation == 'sigmoid': + self.activation = nn.Sigmoid() + elif activation == 'tanh': + self.activation = nn.Tanh() + elif activation == 'relu': + self.activation = nn.ReLU() + else: + raise ValueError(f"Not supported activation `{activation}`.") + + self.lora = nn.Sequential( + nn.Linear(input_dim, low_rank_dim, bias=False), + self.activation, + nn.Linear(low_rank_dim, output_dim, bias=bias), + ) + try: + from transformers.modeling_utils import _init_weights + except ImportError: + _init_weights = True + if _init_weights: + self.apply(self._initialize_weights) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}(" + s += f"input_dim={self.input_dim}, low_rank_dim={self.low_rank_dim}, output_dim={self.output_dim}" + if not self.bias: + s += f", bias={self.bias}" + s += ")" + return s + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + + # Initialize weights to zero as in original code + nn.init.zeros_(self.lora[0].weight) + original_dtype = self.lora[2].weight.dtype + shape = self.lora[2].weight.shape + # Convert to float32 for numerical stability in orthogonal init + weight_fp32 = self.lora[2].weight.float() + + # Calculate gain based on dimensions + gain = math.sqrt(shape[1] / shape[0]) if shape[1] > shape[0] else 1 + + # Apply orthogonal initialization with scaling factor 0.1 + nn.init.orthogonal_(weight_fp32, gain=gain * 0.1) + + # Convert back to original dtype + self.lora[2].weight.data.copy_(weight_fp32.to(original_dtype)) + # Set Lora[2] bias to zero + if self.lora[2].bias is not None: + nn.init.zeros_(self.lora[2].bias) + + module._is_hf_initialized = True + + def set_bias_value(self, value): + """Set bias to a specific value (for v0, w0 etc.)""" + if self.bias and self.lora[2].bias is not None: + if isinstance(value, torch.Tensor): + # Handle tensor values + self.lora[2].bias.data.copy_(value.to(self.lora[2].bias.dtype)) + else: + # Handle scalar values + nn.init.constant_(self.lora[2].bias, value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lora(x) + + +class LerpLinear(nn.Module): + + def __init__( + self, + input_dim: int, + output_dim: int, + low_rank_dim: int | None = None, + ): + super().__init__() + + self.input_dim = input_dim + self.output_dim = output_dim + self.low_rank_dim = low_rank_dim + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + if low_rank_dim is None: + self.linear = nn.Linear(input_dim, output_dim, bias=False) + else: + self.linear = LoRA(input_dim, output_dim, low_rank_dim) + self.mu = nn.Parameter(torch.zeros(input_dim)) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}" + if self.low_rank_dim is not None: + s += f", low_rank_dim={self.low_rank_dim}" + s += ")" + return s + + def forward(self, x: torch.Tensor, delta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None) -> torch.Tensor: + if delta is None: + delta = token_shift(x, cu_seqlens) + return self.linear(x + delta * self.mu) + + +class DDLerpLinear(nn.Module): + + def __init__( + self, + input_dim: int, + output_dim: int, + low_rank_dim: int | None = None, + ): + super().__init__() + + self.input_dim = input_dim + self.output_dim = output_dim + self.low_rank_dim = low_rank_dim + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + if low_rank_dim is None: + self.linear = nn.Linear(input_dim, output_dim, bias=False) + else: + self.linear = LoRA(input_dim, output_dim, low_rank_dim) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}" + if self.low_rank_dim is not None: + s += f", low_rank_dim={self.low_rank_dim}" + s += ")" + return s + + def forward(self, x: torch.Tensor, mu: torch.Tensor, + delta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None) -> torch.Tensor: + if delta is None: + delta = token_shift(x, cu_seqlens) + return self.linear(x + delta * mu) diff --git a/fla/layers/rwkv7.py b/fla/layers/rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..caa4fe039cc5bca137136af07606dfd5818d5b36 --- /dev/null +++ b/fla/layers/rwkv7.py @@ -0,0 +1,347 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.layers.rwkv6 import LoRA +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import GroupNorm +from fla.modules.l2norm import l2_norm +from fla.modules.token_shift import token_shift +from fla.ops.rwkv7 import chunk_rwkv7, fused_mul_recurrent_rwkv7 +from fla.ops.rwkv7.fused_addcmul import fused_addcmul_rwkv7 +from fla.ops.rwkv7.fused_k_update import fused_k_rwkv7 +from fla.ops.rwkv7.gate_output_correction import gate_output_correction + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class RWKV7Attention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + head_dim: int | None = 64, + num_heads: int | None = None, + decay_low_rank_dim: int | None = None, + gate_low_rank_dim: int | None = None, + a_low_rank_dim: int | None = None, + v_low_rank_dim: int | None = None, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + fuse_norm: bool = False, + value_dim: int = None, + num_hidden_layers: int = None, + **kwargs, + ) -> RWKV7Attention: + super().__init__() + + self.mode = mode + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + self.hidden_size = hidden_size + + self.key_dim = hidden_size + self.value_dim = value_dim if value_dim is not None else hidden_size + if head_dim is None and num_heads is None: + raise ValueError("Either `head_dim` or `num_heads` must be specified.") + elif head_dim is not None: + self.head_dim = head_dim + self.num_heads = int(hidden_size // head_dim) + elif num_heads is not None: + self.head_dim = int(hidden_size // num_heads) + self.num_heads = num_heads + self.head_v_dim = int(self.value_dim // self.num_heads) + + # Increase lora dimension for headdim>64 + factor = self.head_dim / 64 + if decay_low_rank_dim is None: + decay_low_rank_dim = max(32, int(round((2.5 * (hidden_size**0.5)) * factor / 32) * 32)) + self.decay_low_rank_dim = decay_low_rank_dim + else: + self.decay_low_rank_dim = decay_low_rank_dim + + if gate_low_rank_dim is None: + gate_low_rank_dim = max(32, int(round((5 * (hidden_size**0.5)) / 32) * 32)) + self.gate_low_rank_dim = gate_low_rank_dim + else: + self.gate_low_rank_dim = gate_low_rank_dim + + if a_low_rank_dim is None: + a_low_rank_dim = max(32, int(round((2.5 * (hidden_size**0.5)) * factor / 32) * 32)) + self.a_low_rank_dim = a_low_rank_dim + else: + self.a_low_rank_dim = a_low_rank_dim + + if v_low_rank_dim is None: + v_low_rank_dim = max(32, int(round((1.7 * (hidden_size**0.5)) * factor / 32) * 32)) + self.v_low_rank_dim = v_low_rank_dim + else: + self.v_low_rank_dim = v_low_rank_dim + + self.layer_idx = layer_idx + self.num_hidden_layers = num_hidden_layers + self.fuse_norm = fuse_norm + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + self.x_r = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.x_w = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.x_k = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.x_v = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.x_a = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.x_g = nn.Parameter(torch.zeros(1, 1, hidden_size)) + + self.k_k = nn.Parameter(torch.zeros(self.key_dim)) + self.k_a = nn.Parameter(torch.zeros(self.key_dim)) + self.r_k = nn.Parameter(torch.zeros(self.num_heads, self.head_dim)) + + self.r_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.w_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=decay_low_rank_dim, activation='tanh') + if self.layer_idx != 0: + self.v_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=v_low_rank_dim, activation=None) + self.a_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=a_low_rank_dim, activation=None) + self.g_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=gate_low_rank_dim, activation='sigmoid', bias=False) + + if self.fuse_norm: + self.g_norm = GroupNorm( + num_groups=self.num_heads, + hidden_size=self.value_dim, + elementwise_affine=elementwise_affine, + eps=self.head_dim*norm_eps, + bias=True, + ) + else: + self.g_norm = nn.GroupNorm( + num_groups=self.num_heads, + num_channels=self.value_dim, + eps=self.head_dim*norm_eps, + affine=elementwise_affine, + ) + + try: + from transformers.modeling_utils import _init_weights + except ImportError: + _init_weights = True + if _init_weights: + self.apply(self._initialize_weights) + for name, module in self.named_modules(): + module._in_rwkv_module = True + + warnings.warn( + "According to Bo, you are using a potentially buggy FLA implementation of RWKV. " + "If you plan to report any numbers based on this implementation, we strongly recommend " + "cross-checking with the official repo: https://github.com/BlinkDL/RWKV-LM. " + "Bo may disagree with results reported from this version.", + ) + + @torch.no_grad() + @torch.compiler.disable + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + + # Initialize only when we're processing the RWKV7Attention module itself + if isinstance(module, RWKV7Attention) and self.layer_idx is not None: + ratio_0_to_1 = self.layer_idx / (self.num_hidden_layers - 1) # 0 to 1 + ratio_1_to_almost0 = 1.0 - (self.layer_idx / self.num_hidden_layers) # 1 to ~0 + + # Create position-based initialization tensor + ddd = torch.ones(1, 1, self.hidden_size, device=self.x_r.device) + www = torch.zeros(self.hidden_size, device=self.x_r.device) + zigzag = torch.zeros(self.hidden_size, device=self.x_r.device) + linear = torch.zeros(self.hidden_size, device=self.x_r.device) + for n in range(self.hidden_size): + linear[n] = n / (self.hidden_size-1) - 0.5 + zigzag[n] = ((n % self.head_dim) - ((self.head_dim-1) / 2)) / ((self.head_dim-1) / 2) + zigzag[n] = zigzag[n] * abs(zigzag[n]) + www[n] = -6 + 6 * (n / (self.hidden_size - 1)) ** (1 + 1 * ratio_0_to_1 ** 0.3) + ddd[0, 0, n] = n / self.hidden_size + + # Initialize x_* parameters directly + self.x_r.data = (1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0)).to(self.x_r.dtype) + self.x_w.data = (1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0)).to(self.x_w.dtype) + self.x_k.data = (1.0 - torch.pow(ddd, 0.7 * ratio_1_to_almost0)).to(self.x_k.dtype) + self.x_v.data = (1.0 - torch.pow(ddd, 0.7 * ratio_1_to_almost0)).to(self.x_v.dtype) + self.x_a.data = (1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0)).to(self.x_a.dtype) + self.x_g.data = (1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0)).to(self.x_g.dtype) + + # Initialize k_k, k_a, r_k + nn.init.constant_(self.k_a, 1.02) + nn.init.constant_(self.r_k, -0.04) + self.k_k.data.copy_((torch.zeros(self.hidden_size, device=self.k_k.device) + + 0.71 - linear*0.1).to(self.k_k.dtype)) + # Set specific bias values for LoRA modules + # 0.5 comes from F.softplus + self.w_lora.set_bias_value(www + 0.5 + zigzag*2.5) + self.a_lora.set_bias_value(-0.19 + zigzag*0.3 + linear*0.4) + + # v0 initialization - ones (for non-first layers) + if self.layer_idx != 0: + self.v_lora._initialize_weights(self.v_lora) + self.v_lora.set_bias_value(0.73 - linear*0.4) + + # Initialize GroupNorm + self.g_norm.weight.data[:] = ((self.layer_idx + 1) / self.num_hidden_layers) ** 0.7 + + # Initialize Linear projections + self._orthogonal_init(self.r_proj.weight) + self._orthogonal_init(self.k_proj.weight, gain=0.1) + self._orthogonal_init(self.v_proj.weight) + self.o_proj.weight.data.zero_() + + # Clean up temporary tensors to free memory + del ddd, www, zigzag, linear + + module._is_hf_initialized = True + + @staticmethod + def _orthogonal_init(weight, gain=1.0): + oringinal_dtype = weight.dtype + weight = weight.float() + nn.init.orthogonal_(weight, gain=gain) + weight = weight.to(oringinal_dtype) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + v_first: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + batch_size, seq_len, _ = hidden_states.shape + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + am = attention_mask.narrow(1, attention_mask.size(1) - seq_len, seq_len).unsqueeze(-1) + + last_state = get_layer_cache(self, past_key_values) + + if attention_mask is not None: + hidden_states = hidden_states.mul(am) + + # delta [batch_size, seq_len, hidden_size] + # conv_cache [N, D] + if last_state is None: + conv_cache = None + recurrent_state = None + else: + conv_cache = last_state['conv_state'] + recurrent_state = last_state['recurrent_state'] + + delta, conv_state = token_shift( + hidden_states, cu_seqlens, output_cache=True, cache=conv_cache, + ) + xr, xw, xk, xv, xa, xg = fused_addcmul_rwkv7(hidden_states, delta, self.x_r, self.x_w, + self.x_k, self.x_v, self.x_a, self.x_g) + + r = self.r_proj(xr) + # Using bf16 for LoRA computation is numerically safe here because: + # 1. After sigmoid activation: + # - Max absolute error (vs float32): 0.003 + # - Mean absolute error: 0.0004 + # 2. Subsequent scaling by -0.6065 will further reduce relative error + # (error scales linearly with constant multiplication) + # 3. Final compounded error remains within acceptable bounds for bf16 precision + # Empirical observation confirms bf16 introduces no practical degradation + w = -0.6065306597126334 * self.w_lora(xw).sigmoid() + + k = self.k_proj(xk) + v = self.v_proj(xv) + + if self.layer_idx == 0: + v_first = v + else: + v = torch.lerp(v, v_first, self.v_lora(xv).sigmoid()) + a = self.a_lora(xa).sigmoid() + g = self.g_lora(xg) + + if self.fuse_norm: + kk = l2_norm(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim)) + else: + kk = F.normalize(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim), dim=-1, p=2.0) + + # Prefer addcmul over expanded form for numerical stability in bf16: + # 1. Fused Multiply-Add (FMA) in addcmul reduces intermediate rounding: + # - Single op vs original 3 ops (mul, sub, mul) + # - 1 less intermediate value storage (bf16 write->read overhead) + # 2. Mathematically equivalent to k*(1 + (a-1)*self.k_a) + # but with better precision preservation + # 3. Particularly crucial for bf16 where intermediate values easily lose precision + # 4. Pytorch method: k = k.addcmul(k * (a - 1), self.k_a) + k = fused_k_rwkv7(k, a, self.k_a) + + # dealing with left-padding + if attention_mask is not None: + v = v * am + + r, w, k, a = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_dim), (r, w, k, a)) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + + if self.training or seq_len >= 64: + # if training, use chunk mode no matter how short the sequence is + # launching the triton kernel for just one token will actually be slower + o, recurrent_state = chunk_rwkv7( + r=r, + w=w, + k=k, + v=v, + a=-kk, + b=kk * a, + scale=1., + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + safe_gate=True, + chunk_size=64, + ) + else: + o, recurrent_state = fused_mul_recurrent_rwkv7( + r=r, + w=w, + k=k, + v=v, + kk=kk, + a=a, + scale=1., + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=conv_state, + offset=r.shape[1], + ) + + if self.fuse_norm: + o = self.g_norm(rearrange(o, '... h d -> ... (h d)')) + else: + o = self.g_norm(rearrange(o, 'b t h d -> (b t) (h d)')).view(batch_size, seq_len, -1) + + o = gate_output_correction(o, r, k, self.r_k, v, g) + o = self.o_proj(o) + + return o, None, past_key_values, v_first diff --git a/fla/layers/simple_gla.py b/fla/layers/simple_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..c04de96e99657ec8edc681358743125054cab3fb --- /dev/null +++ b/fla/layers/simple_gla.py @@ -0,0 +1,273 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.layers.utils import get_layer_cache, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.activations import ACT2FN +from fla.ops.simple_gla import chunk_simple_gla, fused_recurrent_simple_gla + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class SimpleGatedLinearAttention(nn.Module): + r""" + The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635). # noqa + This layer calls the simplified GLA kernel in which the gating is head-wise instead of elementwise. + + Args: + mode (str, Optional): + Which GLA kernel to use. + Currently available: `chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + gate_logit_normalizer (int, Optional): + The normalizer for the gate logits, appied after `logsigmoid`. Default: 16. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1., + expand_v: float = 1., + num_heads: int = 4, + num_kv_heads: int | None = None, + feature_map: str | None = None, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + gate_fn: str = 'swish', + elementwise_affine: bool | None = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 16, + fuse_norm: bool = True, + layer_idx: int = None, + ) -> SimpleGatedLinearAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.layer_idx = layer_idx + + assert mode in ['chunk', "fused_recurrent"], f"Not supported mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim_per_group, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + + self.gk_proj = nn.Linear(hidden_size, self.num_heads) + + if gate_fn == 'swish' and fuse_norm: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps, + dtype=torch.float32 + ) + self.gate_fn = ACT2FN[gate_fn] + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.gate_logit_normalizer = gate_logit_normalizer + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + gk = self.gk_proj(hidden_states) + + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads) + if self.num_kv_groups > 1: + k, v = (repeat(x, '... (h d) -> ... (h g) d', h=self.num_kv_heads, g=self.num_kv_groups) for x in (k, v)) + else: + k, v = (rearrange(x, '... (h d) -> ... h d', h=self.num_kv_heads) for x in (k, v)) + gk = F.logsigmoid(gk) / self.gate_logit_normalizer + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_simple_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q.shape[1], + ) + + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, 'b t (h d) -> b t h d', h=self.num_heads) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, 'b t h d -> b t (h d)') + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = o * self.gate_fn(g) + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_v_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/utils.py b/fla/layers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2fea6fea69bc4cd35d3bb6ec8fa68d2fa280f41b --- /dev/null +++ b/fla/layers/utils.py @@ -0,0 +1,218 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Code is adapted from flash-attn.bert_padding.py + + +import torch +from einops import rearrange, repeat + +from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask +from fla.utils import tensor_cache + +_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided." + + +class IndexFirstAxis(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, indices): + ctx.save_for_backward(indices) + assert x.ndim >= 2 + ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:] + second_dim = other_shape.numel() + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + # return x[indices] + return torch.gather( + rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim), + ).reshape(-1, *other_shape) + + @staticmethod + def backward(ctx, do): + (indices,) = ctx.saved_tensors + assert do.ndim >= 2 + other_shape = do.shape[1:] + do = rearrange(do, "b ... -> b (...)") + dx = torch.zeros( + [ctx.first_axis_dim, do.shape[1]], + device=do.device, + dtype=do.dtype, + ) + # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + # dx[indices] = do + dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do) + return dx.reshape(ctx.first_axis_dim, *other_shape), None + + +index_first_axis = IndexFirstAxis.apply + + +class IndexPutFirstAxis(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, indices, first_axis_dim): + ctx.save_for_backward(indices) + assert indices.ndim == 1 + assert x.ndim >= 2 + y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype) + # TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + y[indices] = x + # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x) + return y + + @staticmethod + def backward(ctx, do): + (indices,) = ctx.saved_tensors + # TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing. + dx = do[indices] + # dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1])) + return dx, None, None + + +index_put_first_axis = IndexPutFirstAxis.apply + + +@tensor_cache +def get_unpad_data( + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Retrieves indexing data required to repad unpadded (ragged) tensors. + + Args: + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + indices (`torch.Tensor`): + The indices of non-masked tokens from the flattened input sequence. + cu_seqlens (`torch.Tensor`): + The cumulative sequence lengths, used to index into ragged (unpadded) tensors. + `cu_seqlens` shape is [batch_size + 1]. + max_seqlen_in_batch (`int`): + Maximum sequence length in batch. + """ + lens = prepare_lens_from_mask(attention_mask) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = lens.max().item() + cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) + return indices, cu_seqlens, max_seqlen_in_batch + + +def unpad_input( + q: torch.Tensor, + states: tuple[torch.Tensor], + attention_mask: torch.Tensor, + q_len: int, + keepdim: bool = False, +): + """ + Unpads query, key, and values tensors, using a single dimension for all tokens + even though they belong to different batches. + + + Arguments: + q (`torch.Tensor`): + Query state with padding. Shape: [batch_size, q_len, ...]. + states (`Tuple[torch.Tensor]`): + Attention state with padding. Shape: [batch_size, seq_len, ...]. + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid. + q_len (`int`): + Target length. + keepdim (`bool`): + Whether to keep the batch dimension. Default: `False`. + + Return: + q (`torch.Tensor`): + Query state without padding. + Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...]. + states (`Tuple[torch.Tensor]`): + Attention state without padding. + Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...]. + indices_q (`torch.Tensor`): + The indices of non-masked tokens from the flattened input target sequence. + (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), + used to index into ragged (unpadded) tensors. + `cu_seqlens` shape is [batch_size + 1]. + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence + i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask) + batch_size, seq_len, *_ = states[0].shape + + state = tuple( + index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k) + for s in states + ) + + if q_len == seq_len: + q = index_first_axis(rearrange(q, "b s ... -> (b s) ..."), indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif q_len == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) + indices_q = cu_seqlens_q[:-1] + q = q.squeeze(1) + else: + raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)") + + if keepdim: + q = q.unsqueeze(0) + state = tuple(s.unsqueeze(0) for s in state) + + return ( + q, + state, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.LongTensor, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + """ + Args: + hidden_states ([total_tokens, ...]): + where total_tokens denotes the number of tokens in selected in attention_mask. + indices ([total_tokens]): + the indices that represent the non-masked tokens of the original padded input sequence. + batch_size (int): + batch_size size for the padded sequence. + seq_len (int): + maximum sequence length for the padded sequence. + + Return: + hidden_states of shape [batch_size, seq_len, ...] + """ + output = index_put_first_axis(hidden_states, indices, batch_size * seq_len) + return rearrange(output, "(b s) ... -> b s ...", b=batch_size) + + +def require_cache_layer_idx(module, past_key_values): + layer_idx = getattr(module, "layer_idx", None) + if past_key_values is not None and layer_idx is None: + raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__)) + return layer_idx + + +def get_layer_cache(module, past_key_values): + layer_idx = require_cache_layer_idx(module, past_key_values) + if past_key_values is not None and len(past_key_values) > layer_idx: + return past_key_values[layer_idx] + return None + + +def update_layer_cache(module, past_key_values, **kwargs): + layer_idx = require_cache_layer_idx(module, past_key_values) + if past_key_values is not None: + return past_key_values.update(layer_idx=layer_idx, **kwargs) + return None diff --git a/fla/models/__init__.py b/fla/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..14f1014d3a7e3225f625a5e4a5e3184b60cdeb32 --- /dev/null +++ b/fla/models/__init__.py @@ -0,0 +1,3 @@ +# Keep model package imports lazy. +__all__ = [] + diff --git a/fla/models/abc/__init__.py b/fla/models/abc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d67d04284af7a704a7246969ab03bbe6dc7c966 --- /dev/null +++ b/fla/models/abc/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.abc.configuration_abc import ABCConfig +from fla.models.abc.modeling_abc import ABCForCausalLM, ABCModel + +AutoConfig.register(ABCConfig.model_type, ABCConfig, exist_ok=True) +AutoModel.register(ABCConfig, ABCModel, exist_ok=True) +AutoModelForCausalLM.register(ABCConfig, ABCForCausalLM, exist_ok=True) + + +__all__ = ['ABCConfig', 'ABCForCausalLM', 'ABCModel'] diff --git a/fla/models/abc/configuration_abc.py b/fla/models/abc/configuration_abc.py new file mode 100644 index 0000000000000000000000000000000000000000..5aec1042a9c87e10ae8d101232f0110bb6227df0 --- /dev/null +++ b/fla/models/abc/configuration_abc.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class ABCConfig(PretrainedConfig): + + model_type = 'abc' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + gate_low_rank_dim: int = 16, + clamp_min: float = -32, + clamp_max: float = 32, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 4, + num_slots: int | None = 64, + use_short_conv: bool = False, + conv_size: int = 4, + exapnd_k: float = 0.5, + exapnd_v: float = 1, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_rope: bool = True, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.gate_low_rank_dim = gate_low_rank_dim + self.clamp_min = clamp_min + self.clamp_max = clamp_max + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_slots = num_slots + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.expand_k = exapnd_k + self.expand_v = exapnd_v + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_rope = use_rope + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/abc/modeling_abc.py b/fla/models/abc/modeling_abc.py new file mode 100644 index 0000000000000000000000000000000000000000..a7260d75e58562094806711323be15cc08186c7d --- /dev/null +++ b/fla/models/abc/modeling_abc.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.abc import ABCAttention +from fla.layers.attn import Attention +from fla.models.abc.configuration_abc import ABCConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as ABCMLP +from fla.modules.l2warp import l2_warp + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +class ABCBlock(GradientCheckpointingLayer): + + def __init__(self, config: ABCConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = ABCAttention( + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + num_slots=config.num_slots, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + gate_fn=config.hidden_act, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + use_rope=config.use_rope, + clamp_min=config.clamp_min, + clamp_max=config.clamp_max, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = ABCMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class ABCPreTrainedModel(PreTrainedModel): + + config_class = ABCConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['ABCBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class ABCModel(ABCPreTrainedModel): + + def __init__(self, config: ABCConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([ABCBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`ABCModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class ABCForCausalLM(ABCPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = ABCModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/bitnet/__init__.py b/fla/models/bitnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3aa44a14248cd17224fe2a40403fb0f4cc6952ba --- /dev/null +++ b/fla/models/bitnet/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.bitnet.configuration_bitnet import BitNetConfig +from fla.models.bitnet.modeling_bitnet import BitNetForCausalLM, BitNetModel + +AutoConfig.register(BitNetConfig.model_type, BitNetConfig, exist_ok=True) +AutoModel.register(BitNetConfig, BitNetModel, exist_ok=True) +AutoModelForCausalLM.register(BitNetConfig, BitNetForCausalLM, exist_ok=True) + + +__all__ = ['BitNetConfig', 'BitNetForCausalLM', 'BitNetModel'] diff --git a/fla/models/bitnet/configuration_bitnet.py b/fla/models/bitnet/configuration_bitnet.py new file mode 100644 index 0000000000000000000000000000000000000000..96309da1be0a6c128b428b86448aa6cbde02b448 --- /dev/null +++ b/fla/models/bitnet/configuration_bitnet.py @@ -0,0 +1,81 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class BitNetConfig(PretrainedConfig): + + model_type = 'bitnet' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 32, + num_kv_heads: int | None = None, + window_size: int | None = None, + rope_theta: float | None = 10000., + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/bitnet/modeling_bitnet.py b/fla/models/bitnet/modeling_bitnet.py new file mode 100644 index 0000000000000000000000000000000000000000..8e4c9ee696560959dbe3a9a84c2da1cfd5ec94b3 --- /dev/null +++ b/fla/models/bitnet/modeling_bitnet.py @@ -0,0 +1,395 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.bitattn import BitAttention +from fla.models.bitnet.configuration_bitnet import BitNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules.activations import swiglu +from fla.modules.fused_bitlinear import FusedBitLinear +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class BitNetMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + hidden_ratio: int | None = None, + intermediate_size: int | None = None, + hidden_act: str = 'swish', + fuse_swiglu: bool = True, + ) -> BitNetMLP: + super().__init__() + + self.hidden_size = hidden_size + # the final number of params is `hidden_ratio * hidden_size^2` + # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio` + if hidden_ratio is None: + hidden_ratio = 4 + if intermediate_size is None: + intermediate_size = int(hidden_size * hidden_ratio * 2 / 3) + intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256) + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.fuse_swiglu = fuse_swiglu + + if hidden_act != 'swish': + raise ValueError(f'Unsupported hidden_act: {hidden_act}') + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + def forward( + self, + x: torch.Tensor, + **kwargs: Unpack[Any], + ) -> torch.Tensor: + gate, y = self.gate_proj(x), self.up_proj(x) + return self.down_proj(swiglu(gate, y)) + + +class BitNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: BitNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = BitAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + window_size=config.window_size, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = BitNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + **kwargs: Unpack[Any], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attentions,) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +class BitNetPreTrainedModel(PreTrainedModel): + + config_class = BitNetConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['BitNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = False, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, FusedBitLinear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class BitNetModel(BitNetPreTrainedModel): + + def __init__( + self, + config: BitNetConfig, + ) -> BitNetModel: + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([BitNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + if output_attentions: + warnings.warn( + "`BitNetModel` does not support output attention weights now, so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + # embed positions + hidden_states = inputs_embeds + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + next_cache = None + + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_attns] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class BitNetForCausalLM(BitNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = BitNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + logits = None if self.config.fuse_linear_cross_entropy else self.lm_head(hidden_states[:, -logits_to_keep:]) + + loss = None + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/comba/__init__.py b/fla/models/comba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..965d9d584cf439f6722fa761b4f9975c4de7c31d --- /dev/null +++ b/fla/models/comba/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.comba.configuration_comba import CombaConfig +from fla.models.comba.modeling_comba import CombaForCausalLM, CombaModel + +AutoConfig.register(CombaConfig.model_type, CombaConfig, exist_ok=True) +AutoModel.register(CombaConfig, CombaModel, exist_ok=True) +AutoModelForCausalLM.register(CombaConfig, CombaForCausalLM, exist_ok=True) + +__all__ = ['CombaConfig', 'CombaForCausalLM', 'CombaModel'] diff --git a/fla/models/comba/configuration_comba.py b/fla/models/comba/configuration_comba.py new file mode 100644 index 0000000000000000000000000000000000000000..6f17f65e0e1bae5aba7e87253c32d122d1f6432a --- /dev/null +++ b/fla/models/comba/configuration_comba.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class CombaConfig(PretrainedConfig): + model_type = 'comba' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + conv_size: int = 4, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int | None = None, + expand_v: float = 2.0, + use_output_gate: bool = True, + use_short_conv: bool = True, + use_output_correction: bool = True, + use_inner_decay: bool = True, + correction_factor: float = 1., + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 21, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.expand_v = expand_v + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.use_output_correction = use_output_correction + self.correction_factor = correction_factor + self.use_inner_decay = use_inner_decay + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/comba/modeling_comba.py b/fla/models/comba/modeling_comba.py new file mode 100644 index 0000000000000000000000000000000000000000..a3363ec3b84829c4283138689757aca4d303fc10 --- /dev/null +++ b/fla/models/comba/modeling_comba.py @@ -0,0 +1,380 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.comba import Comba +from fla.models.comba.configuration_comba import CombaConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as CombaMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class CombaBlock(GradientCheckpointingLayer): + + def __init__(self, config: CombaConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = Comba( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_output_gate=config.use_output_gate, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = CombaMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class CombaPreTrainedModel(PreTrainedModel): + + config_class = CombaConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['CombaBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, Comba) and next(module.parameters()).device.type != 'meta': + with torch.no_grad(): + if not getattr(module.A_log, '_is_hf_initialized', False): + module.A_log.copy_(nn.init.uniform_(module.A_log, a=0, b=16).log()) + module.A_log._no_weight_decay = True + if not getattr(module.dt_bias, '_is_hf_initialized', False): + dt = torch.exp( + nn.init.uniform_(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001), + ).clamp(min=1e-4) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.dt_bias.copy_(inv_dt) + module.dt_bias._no_weight_decay = True + + elif isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class CombaModel(CombaPreTrainedModel): + + def __init__(self, config: CombaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([CombaBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`CombaModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class CombaForCausalLM(CombaPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = CombaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/delta_net/__init__.py b/fla/models/delta_net/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b43ba34156b3b73d78018e0b248ff277b1219a86 --- /dev/null +++ b/fla/models/delta_net/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.delta_net.configuration_delta_net import DeltaNetConfig +from fla.models.delta_net.modeling_delta_net import DeltaNetForCausalLM, DeltaNetModel + +AutoConfig.register(DeltaNetConfig.model_type, DeltaNetConfig, exist_ok=True) +AutoModel.register(DeltaNetConfig, DeltaNetModel, exist_ok=True) +AutoModelForCausalLM.register(DeltaNetConfig, DeltaNetForCausalLM, exist_ok=True) + +__all__ = ['DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel'] diff --git a/fla/models/delta_net/configuration_delta_net.py b/fla/models/delta_net/configuration_delta_net.py new file mode 100644 index 0000000000000000000000000000000000000000..e432338b6059cc16aa6e01ec37b48f07fd3007c1 --- /dev/null +++ b/fla/models/delta_net/configuration_delta_net.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class DeltaNetConfig(PretrainedConfig): + + model_type = 'delta_net' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_k: float = 1.0, + expand_v: float = 1.0, + use_gate: bool = False, + use_short_conv: bool = True, + conv_size: int = 4, + use_beta: bool = True, + use_output_norm: bool = True, + num_heads: int = 16, + qk_norm: str = 'l2', + qk_activation: str = 'silu', + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.use_beta = use_beta + self.use_output_norm = use_output_norm + self.num_heads = num_heads + self.qk_norm = qk_norm + self.qk_activation = qk_activation + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/delta_net/modeling_delta_net.py b/fla/models/delta_net/modeling_delta_net.py new file mode 100644 index 0000000000000000000000000000000000000000..e6cbaf0eb3c62feae73d1fbc026c16319f2f2520 --- /dev/null +++ b/fla/models/delta_net/modeling_delta_net.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.delta_net import DeltaNet +from fla.models.delta_net.configuration_delta_net import DeltaNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as DeltaNetMLP +from fla.modules.l2warp import l2_warp + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +class DeltaNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: DeltaNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = DeltaNet( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + use_gate=config.use_gate, + use_beta=config.use_beta, + use_short_conv=config.use_short_conv, + use_output_norm=config.use_output_norm, + conv_size=config.conv_size, + qk_norm=config.qk_norm, + qk_activation=config.qk_activation, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = DeltaNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class DeltaNetPreTrainedModel(PreTrainedModel): + + config_class = DeltaNetConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['DeltaNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class DeltaNetModel(DeltaNetPreTrainedModel): + + def __init__(self, config: DeltaNetConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([DeltaNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`DeltaNetModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class DeltaNetForCausalLM(DeltaNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = DeltaNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/deltaformer/__init__.py b/fla/models/deltaformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f8b299b7578fa156cec6798eafe117aa404d5815 --- /dev/null +++ b/fla/models/deltaformer/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.deltaformer.configuration_deltaformer import DeltaFormerConfig +from fla.models.deltaformer.modeling_deltaformer import DeltaFormerForCausalLM, DeltaFormerModel + +AutoConfig.register(DeltaFormerConfig.model_type, DeltaFormerConfig, exist_ok=True) +AutoModel.register(DeltaFormerConfig, DeltaFormerModel, exist_ok=True) +AutoModelForCausalLM.register(DeltaFormerConfig, DeltaFormerForCausalLM, exist_ok=True) + +__all__ = ['DeltaFormerConfig', 'DeltaFormerForCausalLM', 'DeltaFormerModel'] diff --git a/fla/models/deltaformer/configuration_deltaformer.py b/fla/models/deltaformer/configuration_deltaformer.py new file mode 100644 index 0000000000000000000000000000000000000000..d51dbc3f29d7bd47ac6bb61182f4f549741761bf --- /dev/null +++ b/fla/models/deltaformer/configuration_deltaformer.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class DeltaFormerConfig(PretrainedConfig): + model_type = 'deltaformer' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 8, + num_kv_heads: int | None = None, + attn_mode: str = "chunk", + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + qkv_bias: bool = False, + qk_norm: bool = False, + rope_theta: float = 10000., + rope_max_position_embeddings: int | None = None, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + output_attentions: bool = False, + output_hidden_states: bool = False, + **kwargs, + ): + self.hidden_size = hidden_size + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.attn_mode = attn_mode + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + self.rope_theta = rope_theta + self.rope_max_position_embeddings = rope_max_position_embeddings + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + self.output_attentions = output_attentions + self.output_hidden_states = output_hidden_states + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/deltaformer/modeling_deltaformer.py b/fla/models/deltaformer/modeling_deltaformer.py new file mode 100644 index 0000000000000000000000000000000000000000..48a0a69613c9d810265875534294c1112697f469 --- /dev/null +++ b/fla/models/deltaformer/modeling_deltaformer.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.deltaformer import DeltaFormerAttention +from fla.models.deltaformer.configuration_deltaformer import DeltaFormerConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as DeltaFormerMLP +from fla.modules.l2warp import l2_warp + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +class DeltaFormerBlock(GradientCheckpointingLayer): + + def __init__(self, config: DeltaFormerConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = DeltaFormerAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + qkv_bias=config.qkv_bias, + qk_norm=config.qk_norm, + rope_theta=config.rope_theta, + max_position_embeddings=config.rope_max_position_embeddings, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = DeltaFormerMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, _, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + outputs = (hidden_states, None, past_key_values) + return outputs + + +class DeltaFormerPreTrainedModel(PreTrainedModel): + + config_class = DeltaFormerConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['DeltaFormerBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + +class DeltaFormerModel(DeltaFormerPreTrainedModel): + + def __init__(self, config: DeltaFormerConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([DeltaFormerBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn( + "`DeltaFormerModel` does not support output attention weights now, " + "so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class DeltaFormerForCausalLM(DeltaFormerPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: DeltaFormerConfig): + super().__init__(config) + self.model = DeltaFormerModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + def tie_weights(self, *args, **kwargs): + """Tie weights for the model. Accepts any arguments for transformers version compatibility.""" + # Use _tie_or_clone_weights if available (older transformers), otherwise rely on parent class + if hasattr(self, '_tie_or_clone_weights'): + self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings()) + else: + # For newer transformers, rely on _tied_weights_keys + super().tie_weights(*args, **kwargs) + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + # For fused linear cross-entropy we do not materialize logits for the full sequence + logits = None if self.config.fuse_linear_cross_entropy else self.lm_head(hidden_states[:, -logits_to_keep:]) + + loss = None + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/forgetting_transformer/__init__.py b/fla/models/forgetting_transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d9c92eb694b5ae579f4ac19424e7294dcbd34b28 --- /dev/null +++ b/fla/models/forgetting_transformer/__init__.py @@ -0,0 +1,15 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.forgetting_transformer.configuration_forgetting_transformer import ForgettingTransformerConfig +from fla.models.forgetting_transformer.modeling_forgetting_transformer import ( + ForgettingTransformerForCausalLM, + ForgettingTransformerModel, +) + +AutoConfig.register(ForgettingTransformerConfig.model_type, ForgettingTransformerConfig, exist_ok=True) +AutoModel.register(ForgettingTransformerConfig, ForgettingTransformerModel, exist_ok=True) +AutoModelForCausalLM.register(ForgettingTransformerConfig, ForgettingTransformerForCausalLM, exist_ok=True) + + +__all__ = ['ForgettingTransformerConfig', 'ForgettingTransformerForCausalLM', 'ForgettingTransformerModel'] diff --git a/fla/models/forgetting_transformer/configuration_forgetting_transformer.py b/fla/models/forgetting_transformer/configuration_forgetting_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..82ce16cb9e8097bb07f1aaaf1260d1dadcd633f9 --- /dev/null +++ b/fla/models/forgetting_transformer/configuration_forgetting_transformer.py @@ -0,0 +1,82 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class ForgettingTransformerConfig(PretrainedConfig): + + model_type = 'forgetting_transformer' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 32, + num_kv_heads: int | None = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: int | None = None, + use_output_gate: bool = False, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + self.window_size = window_size + self.use_output_gate = use_output_gate + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/forgetting_transformer/modeling_forgetting_transformer.py b/fla/models/forgetting_transformer/modeling_forgetting_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..daa673986856113550fd38fcd83401fe87485e17 --- /dev/null +++ b/fla/models/forgetting_transformer/modeling_forgetting_transformer.py @@ -0,0 +1,358 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.forgetting_attn import ForgettingAttention +from fla.models.forgetting_transformer.configuration_forgetting_transformer import ForgettingTransformerConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as ForgettingTransformerMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class ForgettingTransformerBlock(GradientCheckpointingLayer): + + def __init__(self, config: ForgettingTransformerConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = ForgettingAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + qkv_bias=config.qkv_bias, + qk_norm=config.qk_norm, + window_size=config.window_size, + use_output_gate=config.use_output_gate, + layer_idx=layer_idx, + ) + + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = ForgettingTransformerMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + **kwargs: Unpack[Any], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attentions,) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +class ForgettingTransformerPreTrainedModel(PreTrainedModel): + + config_class = ForgettingTransformerConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['ForgettingTransformerBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = False, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per ForgettingTransformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class ForgettingTransformerModel(ForgettingTransformerPreTrainedModel): + + def __init__( + self, + config: ForgettingTransformerConfig, + ) -> ForgettingTransformerModel: + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([ + ForgettingTransformerBlock(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + if output_attentions: + warnings.warn( + "`ForgettingTransformerModel` does not support output attention weights now, " + "so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + # embed positions + hidden_states = inputs_embeds + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + next_cache = None + + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_attns] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class ForgettingTransformerForCausalLM(ForgettingTransformerPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = ForgettingTransformerModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + logits = None if self.config.fuse_linear_cross_entropy else self.lm_head(hidden_states[:, -logits_to_keep:]) + + loss = None + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + # Enable model parallelism + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/gated_deltanet/__init__.py b/fla/models/gated_deltanet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..30666194ff07f6041b60cf41ebf269811fa3267c --- /dev/null +++ b/fla/models/gated_deltanet/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.gated_deltanet.configuration_gated_deltanet import GatedDeltaNetConfig +from fla.models.gated_deltanet.modeling_gated_deltanet import GatedDeltaNetForCausalLM, GatedDeltaNetModel + +AutoConfig.register(GatedDeltaNetConfig.model_type, GatedDeltaNetConfig, exist_ok=True) +AutoModel.register(GatedDeltaNetConfig, GatedDeltaNetModel, exist_ok=True) +AutoModelForCausalLM.register(GatedDeltaNetConfig, GatedDeltaNetForCausalLM, exist_ok=True) + +__all__ = ['GatedDeltaNetConfig', 'GatedDeltaNetForCausalLM', 'GatedDeltaNetModel'] diff --git a/fla/models/gated_deltanet/configuration_gated_deltanet.py b/fla/models/gated_deltanet/configuration_gated_deltanet.py new file mode 100644 index 0000000000000000000000000000000000000000..653f80191efe3fc36c9fbf37babc1b93ed3ff48c --- /dev/null +++ b/fla/models/gated_deltanet/configuration_gated_deltanet.py @@ -0,0 +1,101 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class GatedDeltaNetConfig(PretrainedConfig): + model_type = 'gated_deltanet' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_v: float = 2.0, + use_gate: bool = True, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int | None = None, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 21, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + self.allow_neg_eigval = allow_neg_eigval + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/gated_deltanet/modeling_gated_deltanet.py b/fla/models/gated_deltanet/modeling_gated_deltanet.py new file mode 100644 index 0000000000000000000000000000000000000000..a4823d4c94c6f0ca51b4dcd41d806fc8d00a86d8 --- /dev/null +++ b/fla/models/gated_deltanet/modeling_gated_deltanet.py @@ -0,0 +1,381 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.gated_deltanet import GatedDeltaNet +from fla.models.gated_deltanet.configuration_gated_deltanet import GatedDeltaNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as GatedDeltaNetMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class GatedDeltaNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: GatedDeltaNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = GatedDeltaNet( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_gate=config.use_gate, + use_short_conv=config.use_short_conv, + allow_neg_eigval=config.allow_neg_eigval, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = GatedDeltaNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class GatedDeltaNetPreTrainedModel(PreTrainedModel): + + config_class = GatedDeltaNetConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['GatedDeltaNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, GatedDeltaNet) and next(module.parameters()).device.type != 'meta': + with torch.no_grad(): + if not getattr(module.A_log, '_is_hf_initialized', False): + module.A_log.copy_(nn.init.uniform_(module.A_log, a=0, b=16).log()) + module.A_log._no_weight_decay = True + if not getattr(module.dt_bias, '_is_hf_initialized', False): + dt = torch.exp( + nn.init.uniform_(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001), + ).clamp(min=1e-4) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.dt_bias.copy_(inv_dt) + module.dt_bias._no_weight_decay = True + + elif isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class GatedDeltaNetModel(GatedDeltaNetPreTrainedModel): + + def __init__(self, config: GatedDeltaNetConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([GatedDeltaNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`GatedDeltaNetModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class GatedDeltaNetForCausalLM(GatedDeltaNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = GatedDeltaNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/gated_deltaproduct/__init__.py b/fla/models/gated_deltaproduct/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5c2e1a9f29bd23d7d39b4698e0bec0d9b1115b00 --- /dev/null +++ b/fla/models/gated_deltaproduct/__init__.py @@ -0,0 +1,14 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.gated_deltaproduct.configuration_gated_deltaproduct import GatedDeltaProductConfig +from fla.models.gated_deltaproduct.modeling_gated_deltaproduct import GatedDeltaProductForCausalLM, GatedDeltaProductModel + +AutoConfig.register(GatedDeltaProductConfig.model_type, GatedDeltaProductConfig, exist_ok=True) +AutoModel.register(GatedDeltaProductConfig, GatedDeltaProductModel, exist_ok=True) +AutoModelForCausalLM.register(GatedDeltaProductConfig, GatedDeltaProductForCausalLM, exist_ok=True) + +__all__ = [ + "GatedDeltaProductConfig", + "GatedDeltaProductForCausalLM", + "GatedDeltaProductModel", +] diff --git a/fla/models/gated_deltaproduct/configuration_gated_deltaproduct.py b/fla/models/gated_deltaproduct/configuration_gated_deltaproduct.py new file mode 100644 index 0000000000000000000000000000000000000000..fe40abfb8d5f6628065aa73159252959d96fc4ab --- /dev/null +++ b/fla/models/gated_deltaproduct/configuration_gated_deltaproduct.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class GatedDeltaProductConfig(PretrainedConfig): + model_type = 'gated_deltaproduct' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + conv_size: int = 4, + head_dim: int = 256, + num_heads: int = 6, + hidden_size: int = 2048, + expand_v: float = 2.0, + use_output_gate: bool = True, + use_short_conv: bool = True, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 21, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + use_forget_gate: bool = False, + allow_neg_eigval: bool = False, + num_householder: int = 1, + **kwargs, + ): + self.attn_mode = attn_mode + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + # DeltaProduct specific + self.allow_neg_eigval = allow_neg_eigval + self.num_householder = num_householder + self.use_forget_gate = use_forget_gate + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/gated_deltaproduct/modeling_gated_deltaproduct.py b/fla/models/gated_deltaproduct/modeling_gated_deltaproduct.py new file mode 100644 index 0000000000000000000000000000000000000000..fa0adab7a8ae067fde351db5612ebced4e312f85 --- /dev/null +++ b/fla/models/gated_deltaproduct/modeling_gated_deltaproduct.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.gated_deltaproduct import GatedDeltaProduct +from fla.models.gated_deltaproduct.configuration_gated_deltaproduct import GatedDeltaProductConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as GatedDeltaProductMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class GatedDeltaProductBlock(GradientCheckpointingLayer): + + def __init__(self, config: GatedDeltaProductConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = GatedDeltaProduct( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + use_output_gate=config.use_output_gate, + use_forget_gate=config.use_forget_gate, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + allow_neg_eigval=config.allow_neg_eigval, + num_householder=config.num_householder, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = GatedDeltaProductMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class GatedDeltaProductPreTrainedModel(PreTrainedModel): + + config_class = GatedDeltaProductConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['GatedDeltaProductBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class GatedDeltaProductModel(GatedDeltaProductPreTrainedModel): + + def __init__(self, config: GatedDeltaProductConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([ + GatedDeltaProductBlock(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`GatedDeltaProductModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class GatedDeltaProductForCausalLM(GatedDeltaProductPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = GatedDeltaProductModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/gla/__init__.py b/fla/models/gla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5a680adbb140621c36fdac9a1935e1ab0970af05 --- /dev/null +++ b/fla/models/gla/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.gla.configuration_gla import GLAConfig +from fla.models.gla.modeling_gla import GLAForCausalLM, GLAModel + +AutoConfig.register(GLAConfig.model_type, GLAConfig, exist_ok=True) +AutoModel.register(GLAConfig, GLAModel, exist_ok=True) +AutoModelForCausalLM.register(GLAConfig, GLAForCausalLM, exist_ok=True) + + +__all__ = ['GLAConfig', 'GLAForCausalLM', 'GLAModel'] diff --git a/fla/models/gla/configuration_gla.py b/fla/models/gla/configuration_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..f71d44f06f4a1e01433bb1277f602c734b8c8827 --- /dev/null +++ b/fla/models/gla/configuration_gla.py @@ -0,0 +1,109 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class GLAConfig(PretrainedConfig): + + model_type = 'gla' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + expand_k: float = 0.5, + expand_v: float = 1., + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 4, + num_kv_heads: int | None = None, + feature_map: str | None = None, + attn_mode: str = "chunk", + use_short_conv: bool = False, + conv_size: int = 4, + use_output_gate: bool = True, + clamp_min: float | None = None, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_gk: bool = True, + use_gv: bool = False, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.feature_map = feature_map + self.attn_mode = attn_mode + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.use_output_gate = use_output_gate + self.clamp_min = clamp_min + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_gk = use_gk + self.use_gv = use_gv + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/gla/modeling_gla.py b/fla/models/gla/modeling_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..3d0385da79783def3e9c18dbb7918f251ee0f6b5 --- /dev/null +++ b/fla/models/gla/modeling_gla.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.gla import GatedLinearAttention +from fla.models.gla.configuration_gla import GLAConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as GLAMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class GLABlock(GradientCheckpointingLayer): + + def __init__(self, config: GLAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = GatedLinearAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + feature_map=config.feature_map, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + use_output_gate=config.use_output_gate, + gate_fn=config.hidden_act, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + clamp_min=config.clamp_min, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = GLAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class GLAPreTrainedModel(PreTrainedModel): + + config_class = GLAConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['GLABlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class GLAModel(GLAPreTrainedModel): + + def __init__(self, config: GLAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([GLABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`GLAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class GLAForCausalLM(GLAPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = GLAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/gsa/__init__.py b/fla/models/gsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..61170f8694ad812c2faa98966caee8007b4faed4 --- /dev/null +++ b/fla/models/gsa/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.gsa.configuration_gsa import GSAConfig +from fla.models.gsa.modeling_gsa import GSAForCausalLM, GSAModel + +AutoConfig.register(GSAConfig.model_type, GSAConfig, exist_ok=True) +AutoModel.register(GSAConfig, GSAModel, exist_ok=True) +AutoModelForCausalLM.register(GSAConfig, GSAForCausalLM, exist_ok=True) + + +__all__ = ['GSAConfig', 'GSAForCausalLM', 'GSAModel'] diff --git a/fla/models/gsa/configuration_gsa.py b/fla/models/gsa/configuration_gsa.py new file mode 100644 index 0000000000000000000000000000000000000000..4d4d893b4bdf81816cffa77eb693c207597ddb1f --- /dev/null +++ b/fla/models/gsa/configuration_gsa.py @@ -0,0 +1,111 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class GSAConfig(PretrainedConfig): + + model_type = 'gsa' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + gate_logit_normalizer: int | None = 8, + clamp_min: float | None = None, + clamp_max: float | None = None, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 4, + num_kv_heads: int | None = None, + num_slots: int | None = 64, + use_short_conv: bool = False, + conv_size: int = 4, + exapnd_k: float = 1, + exapnd_v: float = 1, + feature_map: str = 'swish', + use_output_gate: bool = False, + use_norm: bool = True, + max_position_embeddings: int = 2048, + hidden_act: str = "swish", + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + initializer_range: float = 0.02, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.gate_logit_normalizer = gate_logit_normalizer + self.clamp_min = clamp_min + self.clamp_max = clamp_max + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.num_slots = num_slots + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.expand_k = exapnd_k + self.expand_v = exapnd_v + self.feature_map = feature_map + self.use_output_gate = use_output_gate + self.use_norm = use_norm + self.max_position_embeddings = max_position_embeddings + self.hidden_act = hidden_act + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/gsa/modeling_gsa.py b/fla/models/gsa/modeling_gsa.py new file mode 100644 index 0000000000000000000000000000000000000000..27ee3bac5dee776cb62f6089cec2122b513835ba --- /dev/null +++ b/fla/models/gsa/modeling_gsa.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.gsa import GatedSlotAttention +from fla.models.gsa.configuration_gsa import GSAConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as GSAMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class GSABlock(GradientCheckpointingLayer): + + def __init__(self, config: GSAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = GatedSlotAttention( + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + num_slots=config.num_slots, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + feature_map=config.feature_map, + use_output_gate=config.use_output_gate, + use_norm=config.use_norm, + gate_fn=config.hidden_act, + gate_logit_normalizer=config.gate_logit_normalizer, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = GSAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class GSAPreTrainedModel(PreTrainedModel): + + config_class = GSAConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['GSABlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class GSAModel(GSAPreTrainedModel): + + def __init__(self, config: GSAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([GSABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`GSAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class GSAForCausalLM(GSAPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + + super().__init__(config) + self.model = GSAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + # Enable model parallelism + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/hgrn/__init__.py b/fla/models/hgrn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7ef1f0cd18d1f6aba50f4c8e4410a8a9924d9fe --- /dev/null +++ b/fla/models/hgrn/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.hgrn.configuration_hgrn import HGRNConfig +from fla.models.hgrn.modeling_hgrn import HGRNForCausalLM, HGRNModel + +AutoConfig.register(HGRNConfig.model_type, HGRNConfig, exist_ok=True) +AutoModel.register(HGRNConfig, HGRNModel, exist_ok=True) +AutoModelForCausalLM.register(HGRNConfig, HGRNForCausalLM, exist_ok=True) + + +__all__ = ['HGRNConfig', 'HGRNForCausalLM', 'HGRNModel'] diff --git a/fla/models/hgrn/configuration_hgrn.py b/fla/models/hgrn/configuration_hgrn.py new file mode 100644 index 0000000000000000000000000000000000000000..893e53c63a05054fbfa8421cf90b04a13954e3a5 --- /dev/null +++ b/fla/models/hgrn/configuration_hgrn.py @@ -0,0 +1,95 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class HGRNConfig(PretrainedConfig): + + model_type = 'hgrn' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "fused_recurrent", + hidden_size: int = 2048, + num_hidden_layers: int = 24, + expand_ratio: int | None = 1, + use_short_conv: bool = False, + conv_size: int = 4, + use_lower_bound: bool = True, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.expand_ratio = expand_ratio + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.use_lower_bound = use_lower_bound + self.max_position_embeddings = max_position_embeddings + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.elementwise_affine = elementwise_affine + self.attn = attn + self.norm_eps = norm_eps + self.hidden_act = hidden_act + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/hgrn/modeling_hgrn.py b/fla/models/hgrn/modeling_hgrn.py new file mode 100644 index 0000000000000000000000000000000000000000..52953c70c8f94d3ad8a1c89b90e22ff095cb9adb --- /dev/null +++ b/fla/models/hgrn/modeling_hgrn.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.hgrn import HGRNAttention +from fla.models.hgrn.configuration_hgrn import HGRNConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as HGRNMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class HGRNBlock(GradientCheckpointingLayer): + + def __init__(self, config: HGRNConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = HGRNAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_ratio=config.expand_ratio, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = HGRNMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + lower_bound: torch.Tensor | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + lower_bound=lower_bound, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class HGRNPreTrainedModel(PreTrainedModel): + + config_class = HGRNConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['HGRNBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class HGRNModel(HGRNPreTrainedModel): + + def __init__(self, config: HGRNConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + if config.use_lower_bound: + self.lower_bounds = nn.Parameter(torch.zeros(config.num_hidden_layers, config.hidden_size)) + self.layers = nn.ModuleList([HGRNBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`HGRNModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + + if self.config.use_lower_bound: + lower_bounds = self.lower_bounds.softmax(0, dtype=torch.float) + lower_bounds = lower_bounds.cumsum(0) - lower_bounds[0] + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + lower_bound = lower_bounds[i] if self.config.use_lower_bound else None + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + lower_bound=lower_bound, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class HGRNForCausalLM(HGRNPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = HGRNModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/hgrn2/__init__.py b/fla/models/hgrn2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e323464a32a430abee216a305e7ea51a813dd207 --- /dev/null +++ b/fla/models/hgrn2/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.hgrn2.configuration_hgrn2 import HGRN2Config +from fla.models.hgrn2.modeling_hgrn2 import HGRN2ForCausalLM, HGRN2Model + +AutoConfig.register(HGRN2Config.model_type, HGRN2Config, exist_ok=True) +AutoModel.register(HGRN2Config, HGRN2Model, exist_ok=True) +AutoModelForCausalLM.register(HGRN2Config, HGRN2ForCausalLM, exist_ok=True) + + +__all__ = ['HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model'] diff --git a/fla/models/hgrn2/configuration_hgrn2.py b/fla/models/hgrn2/configuration_hgrn2.py new file mode 100644 index 0000000000000000000000000000000000000000..c05e67adeee6f59550aca70a6590c500dfb83076 --- /dev/null +++ b/fla/models/hgrn2/configuration_hgrn2.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class HGRN2Config(PretrainedConfig): + + model_type = 'hgrn2' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + attn_mode: str = "chunk", + num_heads: int | None = None, + expand_ratio: int | None = 128, + use_short_conv: bool = False, + conv_size: int = 4, + use_lower_bound: bool = True, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.attn_mode = attn_mode + + if expand_ratio is None and num_heads is not None: + expand_ratio = hidden_size // num_heads + elif expand_ratio is not None and num_heads is None: + num_heads = hidden_size // expand_ratio + elif expand_ratio is None and num_heads is None: + raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.") + self.num_heads = num_heads + self.expand_ratio = expand_ratio + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.use_lower_bound = use_lower_bound + self.max_position_embeddings = max_position_embeddings + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/hgrn2/modeling_hgrn2.py b/fla/models/hgrn2/modeling_hgrn2.py new file mode 100644 index 0000000000000000000000000000000000000000..e1330a5522ff017a0a41d2028748d3e41ce30773 --- /dev/null +++ b/fla/models/hgrn2/modeling_hgrn2.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.hgrn2 import HGRN2Attention +from fla.models.hgrn2.configuration_hgrn2 import HGRN2Config +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as HGRN2MLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class HGRN2Block(GradientCheckpointingLayer): + + def __init__(self, config: HGRN2Config, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = HGRN2Attention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + num_heads=config.num_heads, + expand_ratio=config.expand_ratio, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = HGRN2MLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + lower_bound: torch.Tensor | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + lower_bound=lower_bound, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class HGRN2PreTrainedModel(PreTrainedModel): + + config_class = HGRN2Config + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['HGRN2Block'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class HGRN2Model(HGRN2PreTrainedModel): + + def __init__(self, config: HGRN2Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + if config.use_lower_bound: + self.lower_bounds = nn.Parameter(torch.zeros(config.num_hidden_layers, config.hidden_size)) + self.layers = nn.ModuleList([HGRN2Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`HGRN2Model` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + + if self.config.use_lower_bound: + lower_bounds = self.lower_bounds.softmax(0, dtype=torch.float) + lower_bounds = lower_bounds.cumsum(0) - lower_bounds[0] + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + lower_bound = lower_bounds[i] if self.config.use_lower_bound else None + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + lower_bound=lower_bound, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class HGRN2ForCausalLM(HGRN2PreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = HGRN2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/kda/__init__.py b/fla/models/kda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cd5af56117d7604fc3801d2d7f8c471622f80856 --- /dev/null +++ b/fla/models/kda/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.kda.configuration_kda import KDAConfig +from fla.models.kda.modeling_kda import KDAForCausalLM, KDAModel + +AutoConfig.register(KDAConfig.model_type, KDAConfig, exist_ok=True) +AutoModel.register(KDAConfig, KDAModel, exist_ok=True) +AutoModelForCausalLM.register(KDAConfig, KDAForCausalLM, exist_ok=True) + +__all__ = ['KDAConfig', 'KDAForCausalLM', 'KDAModel'] diff --git a/fla/models/kda/configuration_kda.py b/fla/models/kda/configuration_kda.py new file mode 100644 index 0000000000000000000000000000000000000000..89b162925e49380e089cf2ec0a1a71075bde92a3 --- /dev/null +++ b/fla/models/kda/configuration_kda.py @@ -0,0 +1,85 @@ + + +from transformers.configuration_utils import PretrainedConfig + + +class KDAConfig(PretrainedConfig): + model_type = 'kda' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_v: float = 1.0, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + head_dim: int = 128, + num_heads: int = 16, + num_v_heads: int | None = None, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + self.allow_neg_eigval = allow_neg_eigval + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/kda/modeling_kda.py b/fla/models/kda/modeling_kda.py new file mode 100644 index 0000000000000000000000000000000000000000..8a239fce355ef96041f03af51241d197df693289 --- /dev/null +++ b/fla/models/kda/modeling_kda.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.kda import KimiDeltaAttention +from fla.models.kda.configuration_kda import KDAConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as KDAMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class KDABlock(GradientCheckpointingLayer): + def __init__(self, config: KDAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn["layers"]: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn["num_heads"], + num_kv_heads=config.attn["num_kv_heads"], + qkv_bias=config.attn["qkv_bias"], + window_size=config.attn["window_size"], + rope_theta=config.attn["rope_theta"], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = KimiDeltaAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_short_conv=config.use_short_conv, + allow_neg_eigval=config.allow_neg_eigval, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = KDAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class KDAPreTrainedModel(PreTrainedModel): + config_class = KDAConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KDABlock"] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, KimiDeltaAttention) and next(module.parameters()).device.type != "meta": + with torch.no_grad(): + if not getattr(module.A_log, '_is_hf_initialized', False): + module.A_log.copy_(nn.init.uniform_(module.A_log, a=1, b=16).log()) + if not getattr(module.dt_bias, '_is_hf_initialized', False): + dt = torch.exp( + nn.init.uniform_(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001), + ).clamp(min=1e-4) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.dt_bias.copy_(inv_dt) + module.dt_bias._is_hf_initialized = True + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None and not getattr(module.bias, "_is_hf_initialized", False): + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, "reset_parameters"): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, "o_proj"): + p = module.o_proj.weight + elif hasattr(module, "down_proj"): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == "rescale": + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == "zero": + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class KDAModel(KDAPreTrainedModel): + def __init__(self, config: KDAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([KDABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`KDAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class KDAForCausalLM(KDAPreTrainedModel, FLAGenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = KDAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if "past_key_values" in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training and labels is not None + + loss, logits = None, None + if not fuse_linear_and_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, "criterion", None) is None: + if fuse_linear_and_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if fuse_linear_and_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/lightnet/__init__.py b/fla/models/lightnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d3edc95065ce5c16fed4120a1bf9fe34265a08e --- /dev/null +++ b/fla/models/lightnet/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.lightnet.configuration_lightnet import LightNetConfig +from fla.models.lightnet.modeling_lightnet import LightNetForCausalLM, LightNetModel + +AutoConfig.register(LightNetConfig.model_type, LightNetConfig, exist_ok=True) +AutoModel.register(LightNetConfig, LightNetModel, exist_ok=True) +AutoModelForCausalLM.register(LightNetConfig, LightNetForCausalLM, exist_ok=True) + + +__all__ = ['LightNetConfig', 'LightNetForCausalLM', 'LightNetModel'] diff --git a/fla/models/lightnet/configuration_lightnet.py b/fla/models/lightnet/configuration_lightnet.py new file mode 100644 index 0000000000000000000000000000000000000000..8f8c9eefed4c659b6b89f412ea8b8bb66dbcf9f4 --- /dev/null +++ b/fla/models/lightnet/configuration_lightnet.py @@ -0,0 +1,97 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class LightNetConfig(PretrainedConfig): + + model_type = 'lightnet' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + attn_mode: str = "chunk", + num_heads: int | None = None, + expand_ratio: int | None = 128, + use_short_conv: bool = False, + conv_size: int = 4, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + gate_low_rank_dim: int = 128, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.attn_mode = attn_mode + self.num_heads = num_heads + self.expand_ratio = expand_ratio + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.max_position_embeddings = max_position_embeddings + self.gate_low_rank_dim = gate_low_rank_dim + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/lightnet/modeling_lightnet.py b/fla/models/lightnet/modeling_lightnet.py new file mode 100644 index 0000000000000000000000000000000000000000..a577ae69f939a8c740ca0244a5b36c3d5ee499a2 --- /dev/null +++ b/fla/models/lightnet/modeling_lightnet.py @@ -0,0 +1,364 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.lightnet import LightNetAttention +from fla.models.lightnet.configuration_lightnet import LightNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as LightNetMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class LightNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: LightNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = LightNetAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + num_heads=config.num_heads, + expand_ratio=config.expand_ratio, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + gate_low_rank_dim=config.gate_low_rank_dim, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = LightNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class LightNetPreTrainedModel(PreTrainedModel): + + config_class = LightNetConfig + supports_gradient_checkpointing = True + _no_split_modules = ['LightNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class LightNetModel(LightNetPreTrainedModel): + + def __init__(self, config: LightNetConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([LightNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`LightNetModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class LightNetForCausalLM(LightNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = LightNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/linear_attn/__init__.py b/fla/models/linear_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f54c0f3c49bb94a1260090c7ca73f104a1ecb91f --- /dev/null +++ b/fla/models/linear_attn/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.linear_attn.configuration_linear_attn import LinearAttentionConfig +from fla.models.linear_attn.modeling_linear_attn import LinearAttentionForCausalLM, LinearAttentionModel + +AutoConfig.register(LinearAttentionConfig.model_type, LinearAttentionConfig, exist_ok=True) +AutoModel.register(LinearAttentionConfig, LinearAttentionModel, exist_ok=True) +AutoModelForCausalLM.register(LinearAttentionConfig, LinearAttentionForCausalLM, exist_ok=True) + +__all__ = ['LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel'] diff --git a/fla/models/linear_attn/configuration_linear_attn.py b/fla/models/linear_attn/configuration_linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..dd7c5a6cd3c019628f5f018098c605b5d503c0e3 --- /dev/null +++ b/fla/models/linear_attn/configuration_linear_attn.py @@ -0,0 +1,105 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class LinearAttentionConfig(PretrainedConfig): + + model_type = 'linear_attn' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "fused_chunk", + hidden_size: int = 2048, + expand_k: float = 1.0, + expand_v: float = 1.0, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 4, + num_kv_heads: int | None = None, + feature_map: str = "elementwise_product", + tie_feature_map_qk: bool = False, + norm_q: bool = False, + norm_k: bool = False, + norm_feature_map: bool = False, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.feature_map = feature_map + self.tie_feature_map_qk = tie_feature_map_qk + self.norm_q = norm_q + self.norm_k = norm_k + self.norm_feature_map = norm_feature_map + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/linear_attn/modeling_linear_attn.py b/fla/models/linear_attn/modeling_linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..beb32028d557419ce0738e70959c63bfdd2dac08 --- /dev/null +++ b/fla/models/linear_attn/modeling_linear_attn.py @@ -0,0 +1,370 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.linear_attn import LinearAttention +from fla.models.linear_attn.configuration_linear_attn import LinearAttentionConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as LinearAttentionMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class LinearAttentionBlock(GradientCheckpointingLayer): + + def __init__(self, config: LinearAttentionConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = LinearAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + feature_map=config.feature_map, + tie_feature_map_qk=config.tie_feature_map_qk, + norm_q=config.norm_q, + norm_k=config.norm_k, + do_feature_map_norm=config.norm_feature_map, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = LinearAttentionMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class LinearAttentionPreTrainedModel(PreTrainedModel): + + config_class = LinearAttentionConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['LinearAttentionBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class LinearAttentionModel(LinearAttentionPreTrainedModel): + + def __init__(self, config: LinearAttentionConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([LinearAttentionBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn( + "`LinearAttentionModel` does not support output attention weights now, " + "so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + + for i, layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class LinearAttentionForCausalLM(LinearAttentionPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = LinearAttentionModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/log_linear_mamba2/__init__.py b/fla/models/log_linear_mamba2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5943ec50a34f51dabe56b81c85c846cfb459aab6 --- /dev/null +++ b/fla/models/log_linear_mamba2/__init__.py @@ -0,0 +1,11 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.log_linear_mamba2.configuration_log_linear_mamba2 import LogLinearMamba2Config +from fla.models.log_linear_mamba2.modeling_log_linear_mamba2 import LogLinearMamba2ForCausalLM, LogLinearMamba2Model + +AutoConfig.register(LogLinearMamba2Config.model_type, LogLinearMamba2Config, exist_ok=True) +AutoModel.register(LogLinearMamba2Config, LogLinearMamba2Model, exist_ok=True) +AutoModelForCausalLM.register(LogLinearMamba2Config, LogLinearMamba2ForCausalLM, exist_ok=True) + + +__all__ = ['LogLinearMamba2Config', 'LogLinearMamba2ForCausalLM', 'LogLinearMamba2Model'] diff --git a/fla/models/log_linear_mamba2/configuration_log_linear_mamba2.py b/fla/models/log_linear_mamba2/configuration_log_linear_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..dd53f724205eab2799716bb2aa2e1b3764066072 --- /dev/null +++ b/fla/models/log_linear_mamba2/configuration_log_linear_mamba2.py @@ -0,0 +1,18 @@ +from fla.models.mamba2 import Mamba2Config + + +class LogLinearMamba2Config(Mamba2Config): + + model_type = "log_linear_mamba2" + + def __init__( + self, + residual_in_fp32: bool = False, + chunk_size: int = 64, + **kwargs, + ): + super().__init__( + residual_in_fp32=residual_in_fp32, + chunk_size=chunk_size, + **kwargs, + ) diff --git a/fla/models/log_linear_mamba2/modeling_log_linear_mamba2.py b/fla/models/log_linear_mamba2/modeling_log_linear_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..c43b93ca20b872bfa1d96fcc5b410e68ced710a8 --- /dev/null +++ b/fla/models/log_linear_mamba2/modeling_log_linear_mamba2.py @@ -0,0 +1,410 @@ +import math + +import torch +from torch import nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.log_linear_mamba2 import LogLinearMamba2 +from fla.models.log_linear_mamba2.configuration_log_linear_mamba2 import LogLinearMamba2Config +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, GatedMLP, RMSNorm + +logger = logging.get_logger(__name__) + + +class LogLinearMamba2Block(nn.Module): + def __init__(self, config: LogLinearMamba2Config, layer_idx: int) -> None: + super().__init__() + if config.residual_in_fp32: + raise NotImplementedError + self.config = config + self.layer_idx = layer_idx + self.mixer_norm = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mlp_norm = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mixer = LogLinearMamba2( + num_heads=config.num_heads, + head_dim=config.head_dim, + hidden_size=config.hidden_size, + state_size=config.state_size, + expand=config.expand, + n_groups=config.n_groups, + conv_kernel=config.conv_kernel, + use_conv_bias=config.use_conv_bias, + hidden_act=config.hidden_act, + rms_norm=config.rms_norm, + chunk_size=config.chunk_size, + time_step_rank=config.time_step_rank, + time_step_limit=config.time_step_limit, + time_step_min=config.time_step_min, + time_step_max=config.time_step_max, + use_bias=config.use_bias, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + hidden_ratio=4, + intermediate_size=None, + hidden_act="swish", + fuse_swiglu=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ): + residual = hidden_states + hidden_states = self.mixer_norm(hidden_states) + hidden_states, attentions, past_key_values = self.mixer( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm( + hidden_states, residual=residual, prenorm=True, + ) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states, attentions, past_key_values + + +class LogLinearMamba2PreTrainedModel(PreTrainedModel, FLAGenerationMixin): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = LogLinearMamba2Config + base_model_prefix = "backbone" + _no_split_modules = ["LogLinearMamba2Block"] + supports_gradient_checkpointing = True + _supports_cache_class = True + + def _init_weights( + self, + module: nn.Module, + num_residuals_per_layer: int = 2, # HAttention + MLP + ): + """Initialize the weights.""" + if isinstance(module, LogLinearMamba2): + # --- A_log --- + A = torch.arange(1, module.num_heads + 1) + with torch.no_grad(): + if not isinstance(module.A_log, torch.distributed.tensor.DTensor): + module.A_log.copy_(torch.log(A)) + else: + logger.warning_once("`A_log` is a DTensor, skipping initialization") + module.A_log._no_weight_decay = True + + # --- D --- + nn.init.ones_(module.D) + module.D._no_weight_decay = True + + # --- L --- + nn.init.ones_(module.L) + module.L._no_weight_decay = True + + # --- dt_bias --- + dt = torch.exp( + torch.rand(self.config.num_heads) + * ( + math.log(self.config.time_step_max) + - math.log(self.config.time_step_min) + ) + + math.log(self.config.time_step_min), + ).clamp(min=self.config.time_step_floor) + + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + if not isinstance(module.dt_bias, torch.distributed.tensor.DTensor): + module.dt_bias.copy_(inv_dt) + else: + logger.warning_once( + "`dt_bias` is a DTensor, skipping initialization", + ) + module.dt_bias._no_reinit = True + + elif isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + # guard against deprecated behavior + if hasattr(module.bias, "_no_reinit"): + raise ValueError("This is not supposed to happen") + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, "reset_parameters"): + module.reset_parameters() + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, "o_proj"): + # p = module.o_proj.weight + # guard against deprecated behavior + raise ValueError("This is not supposed to happen") + elif hasattr(module, "out_proj"): + p = module.out_proj.weight + elif hasattr(module, "down_proj"): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt( + num_residuals_per_layer * self.config.num_hidden_layers, + ) + + +class LogLinearMamba2Model(LogLinearMamba2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + LogLinearMamba2Block(config, layer_idx=idx) + for idx in range(config.num_hidden_layers) + ], + ) + + self.gradient_checkpointing = False + self.norm_f = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + # Initialize weights and apply final processing + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + for k in state_dict: + if "embedding." in k: + state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) + break + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPast: + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + use_cache = ( + use_cache + if use_cache is not None + else (self.config.use_cache if not self.training else False) + ) + if self.gradient_checkpointing and self.training and (use_cache or past_key_values is not None): + logger.warning_once("Disabling cache because gradient checkpointing replays the forward pass.") + use_cache = False + past_key_values = None + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds", + ) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for mixer_block in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + hidden_states, attentions, past_key_values = self._gradient_checkpointing_func( + mixer_block.__call__, + hidden_states, + attention_mask, + past_key_values, + use_cache, + output_attentions, + ) + else: + hidden_states, attentions, past_key_values = mixer_block( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions and attentions is not None: + all_attns = all_attns + (attentions,) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + i + for i in [hidden_states, past_key_values, all_hidden_states, all_attns] + if i is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns if all_attns else None, + ) + + +class LogLinearMamba2ForCausalLM(LogLinearMamba2PreTrainedModel): + _tied_weights_keys = [] + + def __init__(self, config): + super().__init__(config) + self.backbone = LogLinearMamba2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_input_embeddings(self): + return self.backbone.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.backbone.set_input_embeddings(new_embeddings) + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithPast: + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + hidden_states = outputs[0] + fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training + + loss, logits = None, None + if not fuse_linear_and_cross_entropy or labels is None: + logits = self.lm_head( + hidden_states + if logits_to_keep is None + else hidden_states[:, -logits_to_keep:], + ) + if labels is not None: + if getattr(self, "criterion", None) is None: + if fuse_linear_and_cross_entropy: + criterion = FusedLinearCrossEntropyLoss() + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat( + ( + labels[..., 1:], + torch.full_like(labels[:, :1], criterion.ignore_index), + ), + 1, + ) + if fuse_linear_and_cross_entropy: + loss = criterion( + hidden_states, labels, self.lm_head.weight, self.lm_head.bias, + ) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/mamba/__init__.py b/fla/models/mamba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c0215be29f053ff2bc02bd6b540ff9aba2695257 --- /dev/null +++ b/fla/models/mamba/__init__.py @@ -0,0 +1,11 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.mamba.configuration_mamba import MambaConfig +from fla.models.mamba.modeling_mamba import MambaForCausalLM, MambaModel + +AutoConfig.register(MambaConfig.model_type, MambaConfig, exist_ok=True) +AutoModel.register(MambaConfig, MambaModel, exist_ok=True) +AutoModelForCausalLM.register(MambaConfig, MambaForCausalLM, exist_ok=True) + + +__all__ = ['MambaConfig', 'MambaForCausalLM', 'MambaModel'] diff --git a/fla/models/mamba/configuration_mamba.py b/fla/models/mamba/configuration_mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..c00b16dbc33998add081f3c48327d241aa543ddd --- /dev/null +++ b/fla/models/mamba/configuration_mamba.py @@ -0,0 +1,180 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class MambaConfig(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`MambaModel`]. It is used to instantiate a MAMBA + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the MAMBA + [state-spaces/mamba-2.8b](https://huggingface.co/state-spaces/mamba-2.8b) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*): + Vocabulary size of the Mamba model. + hidden_size (`int`, *optional*): + Dimensionality of the embeddings and hidden states. Default: 2048. + state_size (`int`, *optional*): + Shape of the state space latents. Default: 16. + num_hidden_layers (`int`, *optional*): + Number of hidden layers in the model. Default: 48. + norm_eps (`float`, *optional*): + The epsilon to use in the layer normalization layers. Default: 1e-5. + pad_token_id (`int`, *optional*): + Padding token id. Default: 0. + bos_token_id (`int`, *optional*): + The id of the beginning of sentence token in the vocabulary. Default: 0. + eos_token_id (`int`, *optional*): + The id of the end of sentence token in the vocabulary. Default: 0. + expand (`int`, *optional*): + Expanding factor used to determine the intermediate size. Default: 2. + conv_kernel (`int`, *optional*): + Size of the convolution kernel. Default: 4. + use_bias (`bool`, *optional*): + Whether or not to use bias in ["in_proj", "out_proj"] of the mixer block. Default: `False`. + use_conv_bias (`bool`, *optional*): + Whether or not to use bias in the convolution layer of the mixer block. Default: `True`. + hidden_act (`str`, *optional*): + The non-linear activation function (function or string) in the decoder. Default: `"silu"`. + initializer_range (`float`, *optional*): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. Default: 0.02. + residual_in_fp32 (`bool`, *optional*): + Whether or not residuals should be in `float32`. + If set to `False` residuals will keep the same `dtype` as the rest of the model. Default: `True`. + time_step_rank (`Union[int,str]`, *optional*): + Rank of the the discretization projection matrix. + `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`. Default: `"auto"`. + time_step_scale (`float`, *optional*): + Scale used used to scale `dt_proj.bias`. Default: 1.0. + time_step_min (`float`, *optional*): + Minimum `time_step` used to bound `dt_proj.bias`. Default: 0.001. + time_step_max (`float`, *optional*): + Maximum `time_step` used to bound `dt_proj.bias`. Default: 0.1. + time_step_init_scheme (`float`, *optional*): + Init scheme used for `dt_proj.weight`. Should be one of `["random","uniform"]`. Default: `"random"`. + time_step_floor (`float`, *optional*): + Minimum clamping value of the `dt_proj.bias` layer initialization. Default: 0.0001. + window_size (`int`, *optional*): + The window size used for sliding window attention. Default: 2048. + rescale_prenorm_residual (`bool`, *optional*): + Whether or not to rescale `out_proj` weights when initializing. Default: `False`. + use_cache (`bool`, *optional*): + Whether or not the cache should be used. Default: `True`. + + + Example: + + ```python + >>> from transformers import MambaConfig, MambaModel + + >>> # Initializing a Mamba configuration + >>> configuration = MambaConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = MambaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mamba" + + def __init__( + self, + vocab_size: int = 32000, + hidden_size: int = 2048, + state_size: int = 16, + num_hidden_layers: int = 48, + norm_eps=1e-5, + pad_token_id: int = 0, + bos_token_id: int = 1, + eos_token_id: int = 2, + expand: int = 2, + conv_kernel: int = 4, + use_bias: bool = False, + use_conv_bias: bool = True, + hidden_act: str = "silu", + initializer_range: float = 0.02, + residual_in_fp32: bool = False, + time_step_rank: str = "auto", + time_step_scale: float = 1.0, + time_step_min: float = 0.001, + time_step_max: float = 0.1, + time_step_init_scheme: str = "random", + time_step_floor: float = 1e-4, + rescale_prenorm_residual: bool = False, + use_cache: bool = True, + fuse_norm: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + tie_word_embeddings: bool = False, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.state_size = state_size + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.conv_kernel = conv_kernel + self.expand = expand + self.intermediate_size = int(expand * self.hidden_size) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.use_bias = use_bias + self.use_conv_bias = use_conv_bias + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.time_step_rank = math.ceil(self.hidden_size / 16) if time_step_rank == "auto" else time_step_rank + self.time_step_scale = time_step_scale + self.time_step_min = time_step_min + self.time_step_max = time_step_max + self.time_step_init_scheme = time_step_init_scheme + self.time_step_floor = time_step_floor + self.rescale_prenorm_residual = rescale_prenorm_residual + self.residual_in_fp32 = residual_in_fp32 + self.use_cache = use_cache + self.fuse_norm = fuse_norm + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/mamba/modeling_mamba.py b/fla/models/mamba/modeling_mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..6986f21bc7de8ab6c145ec6fe42009db200e72bf --- /dev/null +++ b/fla/models/mamba/modeling_mamba.py @@ -0,0 +1,331 @@ +# Copyright 2024 state-spaces/mamba org and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +from torch import nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.mamba import Mamba +from fla.models.mamba.configuration_mamba import MambaConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules.l2warp import l2_warp + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class MambaBlock(GradientCheckpointingLayer): + + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.residual_in_fp32 = config.residual_in_fp32 + self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mixer = Mamba( + hidden_size=config.hidden_size, + state_size=config.state_size, + conv_kernel=config.conv_kernel, + intermediate_size=config.intermediate_size, + time_step_rank=config.time_step_rank, + use_bias=config.use_bias, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ): + residual = hidden_states + hidden_states = self.norm(hidden_states) + if self.residual_in_fp32: + residual = residual.to(torch.float32) + + hidden_states, attentions, past_key_values = self.mixer( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + hidden_states = residual + hidden_states + if self.residual_in_fp32: + hidden_states = hidden_states.to(dtype=self.norm.weight.dtype) + return hidden_states, attentions, past_key_values + + +class MambaPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = MambaConfig + base_model_prefix = 'backbone' + _no_split_modules = ['Mamba', 'MambaBlock'] + supports_gradient_checkpointing = True + _supports_cache_class = True + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + if not getattr(module.bias, "_no_reinit", False): + nn.init.zeros_(module.bias) + elif isinstance(module, Mamba) and next(module.parameters()).device.type != 'meta': + # S4D real initialization + A = torch.arange(1, module.ssm_state_size + 1, dtype=torch.float32)[None, :] + A = A.expand(module.intermediate_size, -1).contiguous() + with torch.no_grad(): + module.A_log.copy_(torch.log(A)) + module.A_log._no_weight_decay = True + + nn.init.ones_(module.D) + module.D._no_weight_decay = True + + dt_init_std = self.config.time_step_rank**-0.5 * self.config.time_step_scale + if self.config.time_step_init_scheme == "constant": + nn.init.constant_(module.dt_proj.weight, dt_init_std) + elif self.config.time_step_init_scheme == "random": + nn.init.uniform_(module.dt_proj.weight, -dt_init_std, dt_init_std) + + dt = torch.exp( + torch.rand(self.config.intermediate_size) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min), + ).clamp(min=self.config.time_step_floor) + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + module.dt_proj.bias.data = nn.Parameter(inv_dt.to(module.dt_proj.bias.device)) + module.dt_proj.bias._no_reinit = True + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name in ["out_proj.weight"]: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(self.config.num_hidden_layers) + + +class MambaModel(MambaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([MambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.gradient_checkpointing = False + self.norm_f = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + # Initialize weights and apply final processing + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + for k in state_dict: + if "embedding." in k: + state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) + break + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPast: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one", + ) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for mixer_block in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states, attentions, past_key_values = mixer_block( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions and attentions is not None: + all_attns = all_attns + (attentions,) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns if all_attns else None, + ) + + +class MambaForCausalLM(MambaPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.backbone = MambaModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_input_embeddings(self): + return self.backbone.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.backbone.set_input_embeddings(new_embeddings) + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/mamba2/__init__.py b/fla/models/mamba2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98937da0a71289f596fa6b45721eefb4e7668185 --- /dev/null +++ b/fla/models/mamba2/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.mamba2.configuration_mamba2 import Mamba2Config +from fla.models.mamba2.modeling_mamba2 import Mamba2ForCausalLM, Mamba2Model + +AutoConfig.register(Mamba2Config.model_type, Mamba2Config, exist_ok=True) +AutoModel.register(Mamba2Config, Mamba2Model, exist_ok=True) +AutoModelForCausalLM.register(Mamba2Config, Mamba2ForCausalLM, exist_ok=True) + + +__all__ = ['Mamba2Config', 'Mamba2ForCausalLM', 'Mamba2Model'] diff --git a/fla/models/mamba2/configuration_mamba2.py b/fla/models/mamba2/configuration_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..db61364d93a08ab5ef5e11c244557b407b28cb58 --- /dev/null +++ b/fla/models/mamba2/configuration_mamba2.py @@ -0,0 +1,182 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class Mamba2Config(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`Mamba2Model`]. It is used to instantiate a MAMBA2 + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the MAMBA2 + [state-spaces/mamba2-2.8b](https://huggingface.co/state-spaces/mamba2-2.8b) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + head_dim (`int`, *optional*, defaults to 64): + Dimension of each head. + vocab_size (`int`, *optional*, defaults to 32768): + Vocabulary size of the MAMBA2 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Mamba2Model`]. + hidden_size (`int`, *optional*, defaults to 2048): + Dimensionality of the embeddings and hidden states. + state_size (`int`, *optional*, defaults to 128): shape of the state space latents. + num_hidden_layers (`int`, *optional*, defaults to 48): + Number of hidden layers in the model. + norm_eps (`float`, *optional*, defaults to 1e-05): + The epsilon to use in the layer normalization layers. + pad_token_id (`int`, *optional*, defaults to 0): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + The id of the beginning of sentence token in the vocabulary. + eos_token_id (`int`, *optional*, defaults to 2): + The id of the end of sentence token in the vocabulary. + expand (`int`, *optional*, defaults to 2): Expanding factor used to determine the intermediate size. + conv_kernel (`int`, *optional*, defaults to 4): Size of the convolution kernel. + n_groups (`int`, *optional*, defaults to 1): + Number of groups for the evolution matrices of mamba 2. + use_bias (`bool`, *optional*, defaults to `False`): + Whether or not to use bias in ["in_proj", "out_proj"] of the mixer block + use_conv_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use bias in the convolution layer of the mixer block. + hidden_act (`str`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + residual_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether or not residuals should be in `float32`. + If set to `False` residuals will keep the same `dtype` as the rest of the model + time_step_rank (`Union[int,str]`, *optional*, defaults to `"auto"`): + Rank of the discretization projection matrix. + `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)` + time_step_min (`float`, *optional*, defaults to 0.001): + Minimum `time_step` used to bound `dt_proj.bias`. + time_step_max (`float`, *optional*, defaults to 0.1): + Maximum `time_step` used to bound `dt_proj.bias`. + time_step_floor (`float`, *optional*, defaults to 0.0001): + Minimum clamping value of the `dt_proj.bias` layer initialization. + time_step_limit (`tuple`, *optional*, defaults to `(0.0, inf)`): + Accepted range of time step values. + rescale_prenorm_residual (`bool`, *optional*, defaults to `True`): + Whether or not to rescale `out_proj` weights when initializing. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the cache should be used. + rms_norm (`bool`, *optional*, defaults to `True`): + Whether to use RMS norm or not. + chunk_size (`int`, *optional*, defaults to 256): + Size of the chunks that will comprise the sequence. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie word embeddings or not. + """ + + model_type = "mamba2" + + def __init__( + self, + head_dim: int = 64, + vocab_size: int = 32000, + hidden_size: int = 2048, + state_size: int = 128, + num_hidden_layers: int = 48, + norm_eps: float = 1e-5, + pad_token_id: int = 0, + bos_token_id: int = 1, + eos_token_id: int = 2, + expand: int = 2, + conv_kernel: int = 4, + n_groups: int = 1, + use_bias: bool = False, + use_conv_bias: bool = True, + hidden_act: str = "silu", + initializer_range: float = 0.02, + residual_in_fp32: bool = True, + time_step_rank: str = "auto", + time_step_min: float = 0.001, + time_step_max: float = 0.1, + time_step_floor: float = 1e-4, + time_step_limit=(0.0, float("inf")), + rescale_prenorm_residual: bool = True, + use_cache: bool = True, + rms_norm: bool = True, + chunk_size: int = 256, + fuse_norm: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + tie_word_embeddings: bool = False, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.state_size = state_size + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.conv_kernel = conv_kernel + self.expand = expand + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.use_bias = use_bias + self.use_conv_bias = use_conv_bias + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.time_step_rank = ( + math.ceil(self.hidden_size / 16) + if time_step_rank == "auto" + else time_step_rank + ) + self.time_step_min = time_step_min + self.time_step_max = time_step_max + self.time_step_floor = time_step_floor + self.rescale_prenorm_residual = rescale_prenorm_residual + self.residual_in_fp32 = residual_in_fp32 + self.use_cache = use_cache + self.n_groups = n_groups + self.head_dim = head_dim + self.num_heads = int(self.expand * self.hidden_size / self.head_dim) + self.rms_norm = rms_norm + self.state_size = state_size + self.chunk_size = chunk_size + self.time_step_limit = time_step_limit + self.fuse_norm = fuse_norm + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.tie_word_embeddings = tie_word_embeddings + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/mamba2/modeling_mamba2.py b/fla/models/mamba2/modeling_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..0a2cb9eecf06f0b0227220e7e8ff5e060c9f7e5c --- /dev/null +++ b/fla/models/mamba2/modeling_mamba2.py @@ -0,0 +1,398 @@ +# Copyright 2024 state-spaces/mamba2 org and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +from torch import nn +from torch.distributed._tensor.placement_types import Placement, Replicate +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.mamba2 import Mamba2 +from fla.models.mamba2.configuration_mamba2 import Mamba2Config +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules.l2warp import l2_warp + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + + +logger = logging.get_logger(__name__) + + +def tensor_to_dtensor( + tensor: torch.Tensor, + device_mesh: DeviceMesh, + current_placement: Placement | list[Placement], + desired_placement: Placement | list[Placement] | None = None, + run_check: bool = False, +): + if isinstance(tensor, DTensor): + return tensor + + if isinstance(current_placement, Placement): + current_placement = [current_placement] + + dtensor = DTensor.from_local(tensor, device_mesh=device_mesh, run_check=run_check, placements=current_placement) + + if desired_placement is not None: + if isinstance(desired_placement, Placement): + desired_placement = [desired_placement] + + dtensor = dtensor.redistribute(device_mesh=device_mesh, placements=desired_placement, async_op=True) + + return dtensor + + +class Mamba2Block(GradientCheckpointingLayer): + + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.residual_in_fp32 = config.residual_in_fp32 + self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mixer = Mamba2( + num_heads=config.num_heads, + head_dim=config.head_dim, + hidden_size=config.hidden_size, + state_size=config.state_size, + expand=config.expand, + n_groups=config.n_groups, + conv_kernel=config.conv_kernel, + use_conv_bias=config.use_conv_bias, + hidden_act=config.hidden_act, + rms_norm=config.rms_norm, + chunk_size=config.chunk_size, + time_step_rank=config.time_step_rank, + time_step_limit=config.time_step_limit, + time_step_min=config.time_step_min, + time_step_max=config.time_step_max, + use_bias=config.use_bias, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs, + ): + residual = hidden_states + hidden_states = self.norm(hidden_states) + if self.residual_in_fp32: + residual = residual.to(torch.float32) + + hidden_states, attentions, past_key_values = self.mixer( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + hidden_states = residual + hidden_states + if self.residual_in_fp32: + hidden_states = hidden_states.to(dtype=self.norm.weight.dtype) + return hidden_states, attentions, past_key_values + + +class Mamba2PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = Mamba2Config + base_model_prefix = "backbone" + _no_split_modules = ["Mamba2Block"] + supports_gradient_checkpointing = True + _supports_cache_class = True + + def _init_weights( + self, + module: nn.Module, + num_residuals_per_layer: int = 1, + ): + """Initialize the weights.""" + if isinstance(module, Mamba2) and next(module.parameters()).device.type != 'meta': + + # --- A_log --- + A = torch.empty(module.num_heads, dtype=torch.float32).uniform_(0, 16) + with torch.no_grad(): + A_log = torch.log(A) + if isinstance(module.A_log, DTensor): + A_log = tensor_to_dtensor( + tensor=A_log, + device_mesh=module.A_log.device_mesh, + current_placement=[Replicate()] * len(module.A_log.placements), + desired_placement=module.A_log.placements, + run_check=True, + ) + + module.A_log.copy_(A_log) + + module.A_log._no_weight_decay = True + + # --- D --- + nn.init.ones_(module.D) + module.D._no_weight_decay = True + + # --- dt_bias --- + dt = torch.exp( + torch.rand(self.config.num_heads) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min), + ).clamp(min=self.config.time_step_floor) + + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + if isinstance(module.dt_bias, DTensor): + inv_dt = tensor_to_dtensor( + tensor=inv_dt, + device_mesh=module.dt_bias.device_mesh, + current_placement=[Replicate()] * len(module.dt_bias.placements), + desired_placement=module.dt_bias.placements, + run_check=True, + ) + + module.dt_bias.copy_(inv_dt) + module.dt_bias._no_reinit = True + + elif isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + # guard against deprecated behavior + if hasattr(module.bias, "_no_reinit"): + raise ValueError("This is not supposed to happen") + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + # p = module.o_proj.weight + # guard against deprecated behavior + raise ValueError("This is not supposed to happen") + elif hasattr(module, 'out_proj'): + p = module.out_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class Mamba2Model(Mamba2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([Mamba2Block(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.gradient_checkpointing = False + self.norm_f = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + # Initialize weights and apply final processing + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + for k in state_dict: + if "embedding." in k: + state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) + break + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPast: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for mixer_block in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states, attentions, past_key_values = mixer_block( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns = all_attns + (attentions,) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class Mamba2ForCausalLM(Mamba2PreTrainedModel, FLAGenerationMixin): + _tied_weights_keys = [] + + def __init__(self, config): + super().__init__(config) + self.backbone = Mamba2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_input_embeddings(self): + return self.backbone.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.backbone.set_input_embeddings(new_embeddings) + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs, + ) -> tuple | CausalLMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/mesa_net/__init__.py b/fla/models/mesa_net/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..65907264e54ab47c6c410c4fac3b23587eaf1846 --- /dev/null +++ b/fla/models/mesa_net/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.mesa_net.configuration_mesa_net import MesaNetConfig +from fla.models.mesa_net.modeling_mesa_net import MesaNetForCausalLM, MesaNetModel + +AutoConfig.register(MesaNetConfig.model_type, MesaNetConfig, exist_ok=True) +AutoModel.register(MesaNetConfig, MesaNetModel, exist_ok=True) +AutoModelForCausalLM.register(MesaNetConfig, MesaNetForCausalLM, exist_ok=True) + +__all__ = ['MesaNetConfig', 'MesaNetForCausalLM', 'MesaNetModel'] diff --git a/fla/models/mesa_net/configuration_mesa_net.py b/fla/models/mesa_net/configuration_mesa_net.py new file mode 100644 index 0000000000000000000000000000000000000000..0c9549fd638c4b23a1b363fdec35b9f9f14a96ce --- /dev/null +++ b/fla/models/mesa_net/configuration_mesa_net.py @@ -0,0 +1,101 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class MesaNetConfig(PretrainedConfig): + model_type = 'mesa_net' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + use_output_gate: bool = False, + use_short_conv: bool = True, + conv_size: int = 4, + num_heads: int = 16, + head_dim: int = 128, + lambda_lower_bound: float = 0.25, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + max_cg_step_training: int = 30, + max_cg_step_decoding: int = 30, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.num_heads = num_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + self.lambda_lower_bound = lambda_lower_bound + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + self.max_cg_step_training = max_cg_step_training + self.max_cg_step_decoding = max_cg_step_decoding + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/mesa_net/modeling_mesa_net.py b/fla/models/mesa_net/modeling_mesa_net.py new file mode 100644 index 0000000000000000000000000000000000000000..a5332fc570e148ecc7874d31d77cc141cb361f2f --- /dev/null +++ b/fla/models/mesa_net/modeling_mesa_net.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.mesa_net import MesaNet +from fla.models.mesa_net.configuration_mesa_net import MesaNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as MesaNetMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class MesaNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: MesaNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = MesaNet( + mode=config.attn_mode, + hidden_size=config.hidden_size, + num_heads=config.num_heads, + head_dim=config.head_dim, + use_output_gate=config.use_output_gate, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + lambda_lower_bound=config.lambda_lower_bound, + layer_idx=layer_idx, + max_cg_step_training=config.max_cg_step_training, + max_cg_step_decoding=config.max_cg_step_decoding, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = MesaNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class MesaNetPreTrainedModel(PreTrainedModel): + + config_class = MesaNetConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['MesaNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class MesaNetModel(MesaNetPreTrainedModel): + + def __init__(self, config: MesaNetConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([MesaNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`MesaNetModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class MesaNetForCausalLM(MesaNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = MesaNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/mla/__init__.py b/fla/models/mla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5cfa868ba75c924927d92d591faa2621d4ef8a50 --- /dev/null +++ b/fla/models/mla/__init__.py @@ -0,0 +1,11 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.mla.configuration_mla import MLAConfig +from fla.models.mla.modeling_mla import MLAForCausalLM, MLAModel + +AutoConfig.register(MLAConfig.model_type, MLAConfig, exist_ok=True) +AutoModel.register(MLAConfig, MLAModel, exist_ok=True) +AutoModelForCausalLM.register(MLAConfig, MLAForCausalLM, exist_ok=True) + + +__all__ = ['MLAConfig', 'MLAForCausalLM', 'MLAModel'] diff --git a/fla/models/mla/configuration_mla.py b/fla/models/mla/configuration_mla.py new file mode 100644 index 0000000000000000000000000000000000000000..516749adb1474b932389c725135d0fac7734fc97 --- /dev/null +++ b/fla/models/mla/configuration_mla.py @@ -0,0 +1,96 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class MLAConfig(PretrainedConfig): + + model_type = 'mla' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 16, + q_lora_rank: int | None = 64, + qk_rope_head_dim: int = 64, + kv_lora_rank: int = 512, # following the original Deepseek paper + v_head_dim: int = 128, + qk_nope_head_dim: int = 128, + qk_head_dim: int | None = 192, # qk_nope_head_dim + qk_rope_head_dim + window_size: int | None = None, + rope_theta: float | None = 10000., + max_position_embeddings: int = 2048, + rope_scaling: dict | None = None, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + + # MLA specific args + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.kv_lora_rank = kv_lora_rank + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = qk_head_dim + self.rope_scaling = rope_scaling + + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/mla/modeling_mla.py b/fla/models/mla/modeling_mla.py new file mode 100644 index 0000000000000000000000000000000000000000..98ce5e3865611e052b86494e0a77f5e9b03cba86 --- /dev/null +++ b/fla/models/mla/modeling_mla.py @@ -0,0 +1,355 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.mla import MultiheadLatentAttention +from fla.models.mla.configuration_mla import MLAConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as MLAMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class MLABlock(GradientCheckpointingLayer): + + def __init__(self, config: MLAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = MultiheadLatentAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + q_lora_rank=config.q_lora_rank, + qk_rope_head_dim=config.qk_rope_head_dim, + kv_lora_rank=config.kv_lora_rank, + v_head_dim=config.v_head_dim, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_head_dim=config.qk_head_dim, + window_size=config.window_size, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + rope_scaling=config.rope_scaling, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = MLAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class MLAPreTrainedModel(PreTrainedModel): + + config_class = MLAConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['MLABlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class MLAModel(MLAPreTrainedModel): + + def __init__(self, config: MLAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([MLABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`MLAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class MLAForCausalLM(MLAPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: MLAConfig): + super().__init__(config) + self.model = MLAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/modeling_layers.py b/fla/models/modeling_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..24727191a3c478838df4aadba2b579a58f39d5ca --- /dev/null +++ b/fla/models/modeling_layers.py @@ -0,0 +1,70 @@ + +from functools import partial + +from torch import nn +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class GradientCheckpointingLayer(nn.Module): + """Base class for layers with gradient checkpointing. + + This class enables gradient checkpointing functionality for a layer. + By default, gradient checkpointing is disabled (`gradient_checkpointing = False`). + When `model.set_gradient_checkpointing()` is called, gradient checkpointing is enabled + by setting `gradient_checkpointing = True` and assigning a checkpointing function to `_gradient_checkpointing_func`. + + Important: + + When using gradient checkpointing with `use_reentrant=True`, inputs that require gradients (e.g. hidden states) + must be passed as positional arguments (`*args`) rather than keyword arguments to properly propagate gradients. + + Example: + + ```python + >>> # Correct - hidden_states passed as positional arg + >>> out = self.layer(hidden_states, attention_mask=attention_mask) + + >>> # Incorrect - hidden_states passed as keyword arg + >>> out = self.layer(hidden_states=hidden_states, attention_mask=attention_mask) + ``` + """ + + gradient_checkpointing = False + + def __call__(self, *args, **kwargs): + if self.gradient_checkpointing and self.training: + do_warn = False + layer_name = self.__class__.__name__ + message = f"Caching is incompatible with gradient checkpointing in {layer_name}. Setting" + + if "use_cache" in kwargs and kwargs["use_cache"]: + kwargs["use_cache"] = False + message += " `use_cache=False`," + do_warn = True + + # different names for the same thing in different layers + # TODO cyril: this one without `S` can be removed after deprection cycle + if "past_key_value" in kwargs and kwargs["past_key_value"] is not None: + kwargs["past_key_value"] = None + message += " `past_key_value=None`," + do_warn = True + + if "past_key_values" in kwargs and kwargs["past_key_values"] is not None: + kwargs["past_key_values"] = None + message += " `past_key_values=None`," + do_warn = True + + if "layer_past" in kwargs and kwargs["layer_past"] is not None: + kwargs["layer_past"] = None + message += " `layer_past=None`," + do_warn = True + + # warn if anything was changed + if do_warn: + message = message.rstrip(",") + "." + logger.warning_once(message) + + return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) + return super().__call__(*args, **kwargs) diff --git a/fla/models/mom/__init__.py b/fla/models/mom/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..76f5d081b55c3095b1a2ff3046eca07fe40ced25 --- /dev/null +++ b/fla/models/mom/__init__.py @@ -0,0 +1,11 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.mom.configuration_mom import MomConfig +from fla.models.mom.modeling_mom import MomForCausalLM, MomModel + +AutoConfig.register(MomConfig.model_type, MomConfig, exist_ok=True) +AutoModel.register(MomConfig, MomModel, exist_ok=True) +AutoModelForCausalLM.register(MomConfig, MomForCausalLM, exist_ok=True) + +__all__ = ['MomConfig', 'MomForCausalLM', 'MomModel'] diff --git a/fla/models/mom/configuration_mom.py b/fla/models/mom/configuration_mom.py new file mode 100644 index 0000000000000000000000000000000000000000..e16931d9f500814139a5038e94ef82442fdf4feb --- /dev/null +++ b/fla/models/mom/configuration_mom.py @@ -0,0 +1,99 @@ + + +from transformers.configuration_utils import PretrainedConfig + + +class MomConfig(PretrainedConfig): + model_type = 'mom' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + conv_size: int = 4, + num_heads: int = 4, + head_dim: int = 256, + expand_v: float = 1., + use_output_gate: bool = True, + use_short_conv: bool = True, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + num_memories: int = 4, + topk: int = 2, + capacity: float = 1.0, + use_layer_wise_balance: bool = True, + aux_loss_scale: float = 0.01, + shared_mem: bool = True, + single_kv_proj: bool = False, + mom_backend: str = 'gated_deltanet', + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = head_dim + self.expand_v = expand_v + self.conv_size = conv_size + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.num_memories = num_memories + self.topk = topk + self.capacity = capacity + self.use_layer_wise_balance = use_layer_wise_balance + self.aux_loss_scale = aux_loss_scale + self.shared_mem = shared_mem + self.single_kv_proj = single_kv_proj + self.mom_backend = mom_backend + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.vocab_size = vocab_size + + if self.mom_backend not in ['gated_deltanet']: + raise NotImplementedError(f"The MoM backend {mom_backend} is not currently supported.") + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['window_size'] = attn.get('window_size', None) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/mom/modeling_mom.py b/fla/models/mom/modeling_mom.py new file mode 100644 index 0000000000000000000000000000000000000000..5899831072fb1e53b19bdd05d76547a9e841ad43 --- /dev/null +++ b/fla/models/mom/modeling_mom.py @@ -0,0 +1,471 @@ +from __future__ import annotations + +import math +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging + +from fla.layers import MomAttention +from fla.layers.attn import Attention +from fla.models.mom.configuration_mom import MomConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as MomMLP + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +def load_balancing_loss_func( + gate_logits: torch.Tensor | tuple[torch.Tensor] | None, + num_experts: int | None = None, + top_k=2, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor | int: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0, + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0, + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +class MomBlock(GradientCheckpointingLayer): + + def __init__(self, config: MomConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.attn_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + window_size=config.attn['window_size'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + if config.mom_backend == 'gated_deltanet': + self.attn = MomAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + use_output_gate=config.use_output_gate, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + num_memories=config.num_memories, + topk=config.topk, + capacity=config.capacity, + shared_mem=config.shared_mem, + single_kv_proj=config.single_kv_proj, + ) + else: + raise NotImplementedError(f"The MoM backend {config.mom_backend} is not currently supported.") + self.mlp_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mlp = MomMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + if hasattr(self, 'attn_norm'): + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values, router_logits = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if hasattr(self, 'mlp_norm'): + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values, router_logits) + + return outputs + + +class MomPreTrainedModel(PreTrainedModel): + + config_class = MomConfig + supports_gradient_checkpointing = True + _no_split_modules = ['MomBlock'] + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = False, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name in ["o_proj.weight", "down_proj.weight"]: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +@dataclass +class MomOutputWithPast(BaseModelOutputWithPast): + router_logits: tuple[torch.FloatTensor, ...] | None = None + + +class MomModel(MomPreTrainedModel): + + def __init__(self, config: MomConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([MomBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`MomModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + all_router_logits = () + + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values, router_logits = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + all_router_logits += (router_logits,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return MomOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + router_logits=all_router_logits, + ) + + +@dataclass +class MomCausalLMOutputWithPast(CausalLMOutputWithPast): + aux_loss: torch.FloatTensor | None = None + router_logits: tuple[torch.FloatTensor, ...] | None = None + + +class MomForCausalLM(MomPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = MomModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.num_memories = config.num_memories + self.topk = config.topk + self.aux_loss_scale = config.aux_loss_scale + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + num_logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training + logits = None if fuse_linear_and_cross_entropy else self.lm_head(hidden_states[:, -num_logits_to_keep:]) + + loss = None + aux_loss = None + if labels is not None: + if self.config.fuse_cross_entropy: + if fuse_linear_and_cross_entropy: + loss_fct = FusedLinearCrossEntropyLoss() + else: + loss_fct = FusedCrossEntropyLoss(inplace_backward=True) + else: + loss_fct = nn.CrossEntropyLoss() + # Enable model parallelism + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], loss_fct.ignore_index)), 1) + if fuse_linear_and_cross_entropy: + loss = loss_fct(hidden_states.view(-1, self.config.hidden_size), + labels.view(-1), + self.lm_head.weight, + self.lm_head.bias) + else: + loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) + + aux_loss = load_balancing_loss_func( + outputs.router_logits, + self.num_memories, + self.topk, + attention_mask, + ) + + # print(aux_loss) + + loss += aux_loss.to(loss.device) * self.aux_loss_scale + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return MomCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + aux_loss=aux_loss, + ) diff --git a/fla/models/nsa/__init__.py b/fla/models/nsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c0cef083fbfcbb948337059d111f4b1a7e12ca04 --- /dev/null +++ b/fla/models/nsa/__init__.py @@ -0,0 +1,11 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.nsa.configuration_nsa import NSAConfig +from fla.models.nsa.modeling_nsa import NSAForCausalLM, NSAModel + +AutoConfig.register(NSAConfig.model_type, NSAConfig, exist_ok=True) +AutoModel.register(NSAConfig, NSAModel, exist_ok=True) +AutoModelForCausalLM.register(NSAConfig, NSAForCausalLM, exist_ok=True) + + +__all__ = ['NSAConfig', 'NSAForCausalLM', 'NSAModel'] diff --git a/fla/models/nsa/configuration_nsa.py b/fla/models/nsa/configuration_nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..56c1002b3f1d7d34180d87ee62599ba10882394c --- /dev/null +++ b/fla/models/nsa/configuration_nsa.py @@ -0,0 +1,89 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class NSAConfig(PretrainedConfig): + + model_type = 'nsa' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 64, + num_kv_heads: int = 4, + head_dim: int = 32, + qkv_bias: bool = False, + block_size: int = 64, + block_counts: int | None = 16, + window_size: int | None = 512, + rope_theta: float | None = 10000., + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.qkv_bias = qkv_bias + self.block_size = block_size + self.block_counts = block_counts + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/nsa/modeling_nsa.py b/fla/models/nsa/modeling_nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..a160956bb860b4df2cb7658bc58aea9434fab06f --- /dev/null +++ b/fla/models/nsa/modeling_nsa.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.nsa import NativeSparseAttention +from fla.models.nsa.configuration_nsa import NSAConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as NSAMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class NSABlock(GradientCheckpointingLayer): + + def __init__(self, config: NSAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = NativeSparseAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + qkv_bias=config.qkv_bias, + block_size=config.block_size, + block_counts=config.block_counts, + window_size=config.window_size, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = NSAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class NSAPreTrainedModel(PreTrainedModel): + + config_class = NSAConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['NSABlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class NSAModel(NSAPreTrainedModel): + + def __init__(self, config: NSAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([NSABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`NSAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class NSAForCausalLM(NSAPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = NSAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/path_attn/__init__.py b/fla/models/path_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d641116033b06d5f907c821d5bc6e42d02cd449 --- /dev/null +++ b/fla/models/path_attn/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.path_attn.configuration_path_attention import PaTHAttentionConfig +from fla.models.path_attn.modeling_path_attention import PaTHAttentionForCausalLM, PaTHAttentionModel + +AutoConfig.register(PaTHAttentionConfig.model_type, PaTHAttentionConfig, exist_ok=True) +AutoModel.register(PaTHAttentionConfig, PaTHAttentionModel, exist_ok=True) +AutoModelForCausalLM.register(PaTHAttentionConfig, PaTHAttentionForCausalLM, exist_ok=True) + + +__all__ = ['PaTHAttentionConfig', 'PaTHAttentionForCausalLM', 'PaTHAttentionModel'] diff --git a/fla/models/path_attn/configuration_path_attention.py b/fla/models/path_attn/configuration_path_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..0fad6f5c16e4a00bd1aeaf02f61bba58649a6eb5 --- /dev/null +++ b/fla/models/path_attn/configuration_path_attention.py @@ -0,0 +1,81 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class PaTHAttentionConfig(PretrainedConfig): + + model_type = 'path_attn' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 32, + num_kv_heads: int | None = None, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + use_forget_gate: bool = False, + use_w_shortconv: bool = True, + use_low_rank_w: bool = True, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + self.use_forget_gate = use_forget_gate + self.use_w_shortconv = use_w_shortconv + self.use_low_rank_w = use_low_rank_w + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/path_attn/modeling_path_attention.py b/fla/models/path_attn/modeling_path_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..2662e1c2a60796353e424e83449d873032c4eaac --- /dev/null +++ b/fla/models/path_attn/modeling_path_attention.py @@ -0,0 +1,357 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.path_attn import PaTHAttention +from fla.models.path_attn.configuration_path_attention import PaTHAttentionConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as PaTHAttentionMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class PaTHAttentionBlock(GradientCheckpointingLayer): + + def __init__(self, config: PaTHAttentionConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = PaTHAttention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + use_forget_gate=config.use_forget_gate, + use_w_shortconv=config.use_w_shortconv, + use_low_rank_w=config.use_low_rank_w, + layer_idx=layer_idx, + ) + + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = PaTHAttentionMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + **kwargs: Unpack[Any], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attentions,) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +class PaTHAttentionPreTrainedModel(PreTrainedModel): + + config_class = PaTHAttentionConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['PaTHAttentionBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = False, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per PaTHAttention Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class PaTHAttentionModel(PaTHAttentionPreTrainedModel): + + def __init__( + self, + config: PaTHAttentionConfig, + ) -> PaTHAttentionModel: + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([ + PaTHAttentionBlock(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + if output_attentions: + warnings.warn( + "`PaTHAttentionModel` does not support output attention weights now, " + "so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + # embed positions + hidden_states = inputs_embeds + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + next_cache = None + + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_attns] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class PaTHAttentionForCausalLM(PaTHAttentionPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = PaTHAttentionModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + logits = None if self.config.fuse_linear_cross_entropy else self.lm_head(hidden_states[:, -logits_to_keep:]) + + loss = None + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + # Enable model parallelism + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/quasar/__init__.py b/fla/models/quasar/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1284ad439214720b91446d6d18906db8b62ac369 --- /dev/null +++ b/fla/models/quasar/__init__.py @@ -0,0 +1,13 @@ +from fla.models.quasar.configuration_quasar import QuasarConfig +from fla.models.quasar.modeling_quasar import ( + QuasarForCausalLM, + QuasarModel, + QuasarPreTrainedModel +) + +__all__ = [ + 'QuasarConfig', + 'QuasarForCausalLM', + 'QuasarModel', + 'QuasarPreTrainedModel' +] diff --git a/fla/models/quasar/configuration_quasar.py b/fla/models/quasar/configuration_quasar.py new file mode 100644 index 0000000000000000000000000000000000000000..f4ff9e7b24dc871fbc848737f119a2480b684c7b --- /dev/null +++ b/fla/models/quasar/configuration_quasar.py @@ -0,0 +1,87 @@ + + +from transformers.configuration_utils import PretrainedConfig + + +class QuasarConfig(PretrainedConfig): + model_type = 'quasar' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_v: float = 1.0, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + head_dim: int = 128, + num_heads: int = 16, + num_v_heads: int | None = None, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + use_l2warp: bool = False, + use_nope: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.use_l2warp = use_l2warp + self.use_nope = use_nope + self.vocab_size = vocab_size + self.allow_neg_eigval = allow_neg_eigval + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/quasar/modeling_quasar.py b/fla/models/quasar/modeling_quasar.py new file mode 100644 index 0000000000000000000000000000000000000000..bd72404fb8eb2f7c3a9765515c06ff61e40958da --- /dev/null +++ b/fla/models/quasar/modeling_quasar.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.quasar import QuasarAttention +from fla.models.quasar.configuration_quasar import QuasarConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as QuasarMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class QuasarBlock(GradientCheckpointingLayer): + def __init__(self, config: QuasarConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn["layers"]: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn["num_heads"], + num_kv_heads=config.attn["num_kv_heads"], + qkv_bias=config.attn["qkv_bias"], + window_size=config.attn["window_size"], + rope_theta=config.attn["rope_theta"], + max_position_embeddings=config.max_position_embeddings, + use_nope=config.use_nope, + layer_idx=layer_idx, + ) + else: + self.attn = QuasarAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_short_conv=config.use_short_conv, + allow_neg_eigval=config.allow_neg_eigval, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = QuasarMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class QuasarPreTrainedModel(PreTrainedModel): + config_class = QuasarConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["QuasarBlock"] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, QuasarAttention) and next(module.parameters()).device.type != "meta": + with torch.no_grad(): + if not getattr(module.A_log, '_is_hf_initialized', False): + module.A_log.copy_(nn.init.uniform_(module.A_log, a=1, b=16).log()) + if not getattr(module.dt_bias, '_is_hf_initialized', False): + dt = torch.exp( + nn.init.uniform_(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001), + ).clamp(min=1e-4) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.dt_bias.copy_(inv_dt) + module.dt_bias._is_hf_initialized = True + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None and not getattr(module.bias, "_is_hf_initialized", False): + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, "reset_parameters"): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, "o_proj"): + p = module.o_proj.weight + elif hasattr(module, "down_proj"): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == "rescale": + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == "zero": + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class QuasarModel(QuasarPreTrainedModel): + def __init__(self, config: QuasarConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([QuasarBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`QuasarModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class QuasarForCausalLM(QuasarPreTrainedModel, FLAGenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = QuasarModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if "past_key_values" in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training and labels is not None + + loss, logits = None, None + if not fuse_linear_and_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, "criterion", None) is None: + if fuse_linear_and_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if fuse_linear_and_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/retnet/__init__.py b/fla/models/retnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9727d610b133c651061840e129de20d3284d5b1e --- /dev/null +++ b/fla/models/retnet/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.retnet.configuration_retnet import RetNetConfig +from fla.models.retnet.modeling_retnet import RetNetForCausalLM, RetNetModel + +AutoConfig.register(RetNetConfig.model_type, RetNetConfig, exist_ok=True) +AutoModel.register(RetNetConfig, RetNetModel, exist_ok=True) +AutoModelForCausalLM.register(RetNetConfig, RetNetForCausalLM, exist_ok=True) + + +__all__ = ['RetNetConfig', 'RetNetForCausalLM', 'RetNetModel'] diff --git a/fla/models/retnet/configuration_retnet.py b/fla/models/retnet/configuration_retnet.py new file mode 100644 index 0000000000000000000000000000000000000000..b3d6db39241c1d95e24bec1a95eb7d7cb11104c1 --- /dev/null +++ b/fla/models/retnet/configuration_retnet.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class RetNetConfig(PretrainedConfig): + + model_type = 'retnet' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_k: float = 1.0, + expand_v: float = 2.0, + hidden_ratio: int | None = 2, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 8, + num_kv_heads: int | None = None, + feature_map: str | None = None, + hidden_act: str = "swish", + use_short_conv: bool = False, + conv_size: int = 4, + use_output_gate: bool = True, + max_position_embeddings: int = 2048, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ) -> RetNetConfig: + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.feature_map = feature_map + self.hidden_act = hidden_act + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.use_output_gate = use_output_gate + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/retnet/modeling_retnet.py b/fla/models/retnet/modeling_retnet.py new file mode 100644 index 0000000000000000000000000000000000000000..bbd0fb7fb0e91af958be1df2e85ebc2e8371c2d4 --- /dev/null +++ b/fla/models/retnet/modeling_retnet.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.multiscale_retention import MultiScaleRetention +from fla.models.retnet.configuration_retnet import RetNetConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as RetNetMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class RetNetBlock(GradientCheckpointingLayer): + + def __init__(self, config: RetNetConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = MultiScaleRetention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + feature_map=config.feature_map, + use_output_gate=config.use_output_gate, + gate_fn=config.hidden_act, + elementwise_affine=config.elementwise_affine, + norm_eps=config.norm_eps, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = RetNetMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class RetNetPreTrainedModel(PreTrainedModel): + + config_class = RetNetConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['RetNetBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class RetNetModel(RetNetPreTrainedModel): + + def __init__(self, config: RetNetConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [RetNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)], + ) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn( + "`RetNetModel` does not support output attention weights now, so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class RetNetForCausalLM(RetNetPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = RetNetModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + # Expected exception: "AttributeError: '(object name)' object has no attribute 'past_key_values'" + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/rodimus/__init__.py b/fla/models/rodimus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0aa7da945f9115d28562203a2087f4d181897bf8 --- /dev/null +++ b/fla/models/rodimus/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.rodimus.configuration_rodimus import RodimusConfig +from fla.models.rodimus.modeling_rodimus import RodimusForCausalLM, RodimusModel + +AutoConfig.register(RodimusConfig.model_type, RodimusConfig, exist_ok=True) +AutoModel.register(RodimusConfig, RodimusModel, exist_ok=True) +AutoModelForCausalLM.register(RodimusConfig, RodimusForCausalLM, exist_ok=True) + + +__all__ = ['RodimusConfig', 'RodimusForCausalLM', 'RodimusModel'] diff --git a/fla/models/rodimus/configuration_rodimus.py b/fla/models/rodimus/configuration_rodimus.py new file mode 100644 index 0000000000000000000000000000000000000000..a389f75356ac3c5f3de75e947df83c86a5118e38 --- /dev/null +++ b/fla/models/rodimus/configuration_rodimus.py @@ -0,0 +1,116 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class RodimusConfig(PretrainedConfig): + + model_type = 'rodimus' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + block_type: str = 'rodimus_plus', + hidden_size: int = 2048, + num_hidden_layers: int = 24, + attn_mode: str = "chunk", + residual_in_fp32: bool = True, + block_residual_in_fp32: bool = False, + expand_ratio: int | None = 64, + input_gate_low_rank: float | str | None = 'auto', + use_short_conv: bool = True, + conv_size: int = 4, + hidden_ratio: float | None = 4/3, + intermediate_size: int | None = None, + hidden_act: str = "swish", + max_position_embeddings: int = 2048, + norm_eps: float = 1e-5, + k_norm_eps: float | None = None, + attn: dict | None = None, + ska_attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 126080, + eos_token_id: int = 126081, + tie_word_embeddings: bool = True, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 126464, + **kwargs, + ): + self.block_type = block_type + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.attn_mode = attn_mode + self.residual_in_fp32 = residual_in_fp32 + self.block_residual_in_fp32 = block_residual_in_fp32 + self.expand_ratio = expand_ratio + self.input_gate_low_rank = input_gate_low_rank + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.max_position_embeddings = max_position_embeddings + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.norm_eps = norm_eps + self.k_norm_eps = k_norm_eps + + self.attn = attn + self.ska_attn = ska_attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['qk_norm'] = attn.get('qk_norm', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + if ska_attn is not None: + if not isinstance(ska_attn, dict): + raise ValueError("attn must be a dictionary") + if 'num_heads' not in ska_attn: + raise ValueError("Number of heads must be provided to initialize shared-key attention layers") + ska_attn['qkv_bias'] = ska_attn.get('qkv_bias', False) + ska_attn['qk_norm'] = ska_attn.get('qk_norm', False) + ska_attn['window_size'] = ska_attn.get('window_size', 1024) + ska_attn['rope_theta'] = ska_attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/rodimus/modeling_rodimus.py b/fla/models/rodimus/modeling_rodimus.py new file mode 100644 index 0000000000000000000000000000000000000000..f282a45c1232d56216520343c35bbfed4e86abcd --- /dev/null +++ b/fla/models/rodimus/modeling_rodimus.py @@ -0,0 +1,560 @@ +from __future__ import annotations + +import math +import warnings +from functools import partial +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.rodimus import RodimusAttention, SlidingWindowSharedKeyAttention, align_multiple +from fla.models.rodimus.configuration_rodimus import RodimusConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as RodimusMLP +from fla.modules.l2warp import l2_warp + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class RodimusBlock(GradientCheckpointingLayer): + + def __init__(self, config: RodimusConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + self.block_type = config.block_type + self.block_residual_in_fp32 = config.residual_in_fp32 + self.residual_in_fp32 = config.residual_in_fp32 + self.fuse_norm = config.fuse_norm + + self._is_ori_attn = False + + if config.intermediate_size is None: + intermediate_size = align_multiple(int(config.hidden_ratio * config.hidden_size), 8) + else: + intermediate_size = config.intermediate_size + + mlp_cls = partial( + RodimusMLP, + hidden_size=config.hidden_size, + hidden_ratio=None, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + norm_cls = partial( + RMSNorm if self.fuse_norm else nn.RMSNorm, + config.hidden_size, + eps=config.norm_eps, + ) + + if config.attn is not None and layer_idx in config.attn['layers']: + self._is_ori_attn = True + self.attn_norm = norm_cls() + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + + self.mlp_norm = norm_cls() + self.mlp = mlp_cls() + else: + self.mixer_norm = norm_cls() + self.mixer = RodimusAttention( + block_type=config.block_type, + mode=config.attn_mode, + hidden_size=config.hidden_size, + input_gate_low_rank=config.input_gate_low_rank, + expand_ratio=config.expand_ratio, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + k_norm_eps=config.k_norm_eps, + residual_in_fp32=config.residual_in_fp32, + layer_idx=layer_idx, + ) + + if self.block_type == "rodimus_plus": + self.ska_attn_norm = norm_cls() + self.ska_attn = SlidingWindowSharedKeyAttention( + hidden_size=config.hidden_size, + num_heads=config.ska_attn['num_heads'], + qkv_bias=config.ska_attn['qkv_bias'], + qk_norm=config.ska_attn['qk_norm'], + window_size=config.ska_attn['window_size'], + rope_theta=config.ska_attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + + self.mlp_norm = norm_cls() + self.mlp = mlp_cls() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + residual: torch.Tensor | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + if self.block_residual_in_fp32 and self.layer_idx > 0: + assert residual is not None, 'Residual must be passed in when setting `block_residual_in_fp32=True`' + + if self._is_ori_attn: + if self.block_residual_in_fp32: + hidden_states, residual = self.attn_norm( + hidden_states, + residual=residual, + prenorm=True, + residual_in_fp32=self.residual_in_fp32, + ) + else: + residual = hidden_states.float() if self.residual_in_fp32 else hidden_states + hidden_states = self.attn_norm(hidden_states) + + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.fuse_norm: + hidden_states, residual = self.mlp_norm( + hidden_states, + residual, + prenorm=True, + residual_in_fp32=self.residual_in_fp32, + ) + else: + hidden_states = residual + hidden_states + residual = hidden_states.float() if self.residual_in_fp32 else hidden_states + hidden_states = self.mlp_norm(hidden_states.to(self.mlp_norm.weight.dtype)) + + hidden_states = self.mlp(hidden_states, **kwargs) + else: + if self.block_residual_in_fp32: + hidden_states, residual = self.mixer_norm( + hidden_states, + residual=residual, + prenorm=True, + residual_in_fp32=self.residual_in_fp32, + ) + else: + residual = hidden_states.float() if self.residual_in_fp32 else hidden_states + hidden_states = self.mixer_norm(hidden_states) + + hidden_states, attentions, past_key_values = self.mixer( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if self.block_type == "rodimus_plus": + past_key_values, rodimus_caches = past_key_values + + if self.fuse_norm: + hidden_states, residual = self.ska_attn_norm( + hidden_states, + residual, + prenorm=True, + residual_in_fp32=self.residual_in_fp32, + ) + else: + hidden_states = residual + hidden_states + residual = hidden_states.float() if self.residual_in_fp32 else hidden_states + hidden_states = self.ska_attn_norm(hidden_states.to(dtype=self.ska_attn_norm.weight.dtype)) + + hidden_states, attentions, past_key_values = self.ska_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + rodimus_caches=rodimus_caches, + **kwargs, + ) + + if self.fuse_norm: + hidden_states = self.mlp_norm( + hidden_states, + residual=residual, + prenorm=False, + residual_in_fp32=self.residual_in_fp32, + ) + else: + hidden_states = residual + hidden_states + hidden_states = self.mlp_norm(hidden_states.to(dtype=self.mlp_norm.weight.dtype)) + + hidden_states = self.mlp(hidden_states, **kwargs) + + if self.block_residual_in_fp32: + hidden_states = (hidden_states, residual) + else: + hidden_states = (residual + hidden_states).to(dtype=hidden_states.dtype) + + outputs = (hidden_states, attentions, past_key_values) + return outputs + + +class RodimusPreTrainedModel(PreTrainedModel): + + config_class = RodimusConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['RodimusBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + if self.config.block_type == "rodimus": + self.num_residuals_per_layer = 1 + elif self.config.block_type == "rodimus_plus": + self.num_residuals_per_layer = 3 + else: + raise NotImplementedError() + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + ): + num_residuals_per_layer = self.num_residuals_per_layer + + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + sigmoid_bias_max = 0.999 + sigmoid_bias_min = 0.9 + max_ = 1 - sigmoid_bias_min + min_ = 1 - sigmoid_bias_max + g_gate_bias = torch.exp( + torch.rand(self.config.expand_ratio) * (math.log(max_) - math.log(min_)) + + math.log(min_), + ).clamp(min=1e-4) + g_gate_bias = g_gate_bias + torch.log(-torch.expm1(-g_gate_bias)) + tau_gate_bias = torch.logit(torch.empty((self.config.expand_ratio, )).uniform_(1/16, 0.9)) + + if hasattr(module, 'i_gate_proj'): + nn.init.xavier_uniform_(module.i_gate_proj[0].weight, gain=2 ** -2.5) + nn.init.xavier_uniform_(module.i_gate_proj[1].weight, gain=2 ** -2.5) + nn.init.zeros_(module.i_gate_proj[1].bias) + if hasattr(module, 'g_gate_proj'): + nn.init.xavier_uniform_(module.g_gate_proj.weight, gain=2 ** -2.5) + with torch.no_grad(): + if not isinstance(module.g_gate_proj.bias, DTensor): + module.g_gate_proj.bias.copy_(g_gate_bias) + else: + logger.warning_once("`g_gate_proj.bias` is a DTensor, skipping initialization") + + if hasattr(module, 'tau_gate_proj'): + nn.init.xavier_uniform_(module.tau_gate_proj.weight, gain=2 ** -2.5) + with torch.no_grad(): + if not isinstance(module.tau_gate_proj.bias, DTensor): + module.tau_gate_proj.bias.copy_(tau_gate_bias) + else: + logger.warning_once("`tau_gate_proj.bias` is a DTensor, skipping initialization") + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class RodimusModel(RodimusPreTrainedModel): + + def __init__(self, config: RodimusConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.block_residual_in_fp32 = config.block_residual_in_fp32 + + if config.block_residual_in_fp32: + if not config.residual_in_fp32: + warning_message = ( + "`residual_in_fp32=False` is incompatible with `block_residual_in_fp32=True`. " + "Setting `residual_in_fp32=True`..." + ) + logger.warning_once(warning_message) + config.residual_in_fp32 = True + if not config.fuse_norm: + logger.warning_once( + '`fuse_norm=False` is incompatible with `block_residual_in_fp32=True` Setting `fuse_norm=True`...') + config.fuse_norm = True + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([RodimusBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + + if output_attentions: + warnings.warn("`RodimusModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + residual = None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + residual=residual, + **kwargs, + ) + + if self.block_residual_in_fp32: + hidden_states, residual = hidden_states + else: + residual = None + + if output_attentions: + all_attns += (attentions,) + + if self.block_residual_in_fp32: + hidden_states = self.norm( + hidden_states, + residual=residual, + prenorm=False, + residual_in_fp32=True, + ) + else: + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class RodimusForCausalLM(RodimusPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = RodimusModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/rwkv6/__init__.py b/fla/models/rwkv6/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bde1972830019cfb6c8f7f7f38507f85c939d3df --- /dev/null +++ b/fla/models/rwkv6/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.rwkv6.configuration_rwkv6 import RWKV6Config +from fla.models.rwkv6.modeling_rwkv6 import RWKV6ForCausalLM, RWKV6Model + +AutoConfig.register(RWKV6Config.model_type, RWKV6Config, exist_ok=True) +AutoModel.register(RWKV6Config, RWKV6Model, exist_ok=True) +AutoModelForCausalLM.register(RWKV6Config, RWKV6ForCausalLM, exist_ok=True) + + +__all__ = ['RWKV6Config', 'RWKV6ForCausalLM', 'RWKV6Model'] diff --git a/fla/models/rwkv6/configuration_rwkv6.py b/fla/models/rwkv6/configuration_rwkv6.py new file mode 100644 index 0000000000000000000000000000000000000000..56559b091ed8789da30c4eadf393bd7002b8fe34 --- /dev/null +++ b/fla/models/rwkv6/configuration_rwkv6.py @@ -0,0 +1,96 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class RWKV6Config(PretrainedConfig): + + model_type = 'rwkv6' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_k: float = 0.5, + expand_v: float = 1.0, + hidden_ratio: float | None = 3.5, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + num_heads: int = 4, + proj_low_rank_dim: int = 32, + gate_low_rank_dim: int = 64, + hidden_act: str = "sqrelu", + max_position_embeddings: int = 2048, + norm_first: bool = True, + norm_bias: bool = True, + norm_eps: float = 1e-5, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.norm_first = norm_first + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.proj_low_rank_dim = proj_low_rank_dim + self.gate_low_rank_dim = gate_low_rank_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.norm_bias = norm_bias + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + self.fuse_norm = fuse_norm + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/rwkv6/modeling_rwkv6.py b/fla/models/rwkv6/modeling_rwkv6.py new file mode 100644 index 0000000000000000000000000000000000000000..e157689afe86d14381a52199ecdb3bfb504461b6 --- /dev/null +++ b/fla/models/rwkv6/modeling_rwkv6.py @@ -0,0 +1,447 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.rwkv6 import LerpLinear, RWKV6Attention +from fla.models.rwkv6.configuration_rwkv6 import RWKV6Config +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, LayerNorm +from fla.modules.activations import ACT2FN +from fla.modules.l2warp import l2_warp +from fla.modules.token_shift import token_shift + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class RWKV6FeedForward(nn.Module): + + def __init__( + self, + hidden_size: int, + hidden_ratio: int | None = None, + intermediate_size: int | None = None, + hidden_act: str = 'sqrelu', + layer_idx: int = None, + ) -> RWKV6FeedForward: + super().__init__() + + self.hidden_size = hidden_size + if hidden_ratio is None: + hidden_ratio = 3.5 + if intermediate_size is None: + intermediate_size = int(hidden_size * hidden_ratio) + intermediate_size = 32 * ((intermediate_size + 32 - 1) // 32) + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + + self.key = LerpLinear(hidden_size, intermediate_size) + self.value = nn.Linear(intermediate_size, hidden_size, bias=False) + self.receptance = LerpLinear(hidden_size, hidden_size) + self.act_fn = ACT2FN[hidden_act] + + self.layer_idx = layer_idx + + def forward( + self, + x: torch.Tensor, + attention_mask: torch.Tensor | None = None, + state: Cache | None = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> torch.Tensor: + if attention_mask is not None: + x = x.mul(attention_mask[:, -x.shape[-2]:, None]) + if x.shape[1] == 1 and state is not None and state[self.layer_idx]['ffn_state'] is not None: + shifted = state[self.layer_idx]['ffn_state'].unsqueeze(1) + delta = shifted - x + elif state is not None and state[self.layer_idx]['ffn_state'] is not None: + shifted = self.time_shift(x) + if state is not None and state[self.layer_idx]['ffn_state'] is not None: + shifted[:, 0] = state[self.layer_idx]['ffn_state'] + delta = shifted - x + else: + delta = token_shift(x, cu_seqlens) + key = self.act_fn(self.key(x, delta)) + value = self.value(key) + receptance = self.receptance(x, delta) + + if state is not None: + # no need to update the offset twice + state.update(ffn_state=x[:, -1], layer_idx=self.layer_idx, offset=0) + return receptance.sigmoid() * value, state + + +class RWKV6Block(GradientCheckpointingLayer): + + def __init__(self, config: RWKV6Config, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + if config.norm_first and layer_idx == 0: + self.pre_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + self.attn_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = RWKV6Attention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_k=config.expand_k, + expand_v=config.expand_v, + num_heads=config.num_heads, + proj_low_rank_dim=config.proj_low_rank_dim, + gate_low_rank_dim=config.gate_low_rank_dim, + norm_eps=config.norm_eps, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + ) + self.ffn_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + self.ffn = RWKV6FeedForward( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = self.pre_norm(hidden_states) if hasattr(self, 'pre_norm') else hidden_states + hidden_states = self.attn_norm(residual) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + cu_seqlens=cu_seqlens, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.ffn_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.ffn_norm(hidden_states) + hidden_states, past_key_values = self.ffn( + hidden_states, attention_mask, past_key_values, cu_seqlens, **kwargs, + ) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class RWKV6PreTrainedModel(PreTrainedModel): + + config_class = RWKV6Config + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['RWKV6Block'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = True, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Parameter): + nn.init.normal_(module, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class RWKV6Model(RWKV6PreTrainedModel): + + def __init__(self, config: RWKV6Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([RWKV6Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`RWKV6Model` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + cu_seqlens=cu_seqlens, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class RWKV6ForCausalLM(RWKV6PreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = RWKV6Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/rwkv7/__init__.py b/fla/models/rwkv7/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2b3fff4d18b1ec37fa440a349fdcd74b4202ed52 --- /dev/null +++ b/fla/models/rwkv7/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.rwkv7.configuration_rwkv7 import RWKV7Config +from fla.models.rwkv7.modeling_rwkv7 import RWKV7ForCausalLM, RWKV7Model + +AutoConfig.register(RWKV7Config.model_type, RWKV7Config, exist_ok=True) +AutoModel.register(RWKV7Config, RWKV7Model, exist_ok=True) +AutoModelForCausalLM.register(RWKV7Config, RWKV7ForCausalLM, exist_ok=True) + + +__all__ = ['RWKV7Config', 'RWKV7ForCausalLM', 'RWKV7Model'] diff --git a/fla/models/rwkv7/configuration_rwkv7.py b/fla/models/rwkv7/configuration_rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..c6bd8a429c627f534a0bf589c6b5994e2ee6bb18 --- /dev/null +++ b/fla/models/rwkv7/configuration_rwkv7.py @@ -0,0 +1,119 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class RWKV7Config(PretrainedConfig): + + model_type = 'rwkv7' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + head_dim: int | None = 64, + num_heads: int | None = None, + decay_low_rank_dim: int = 64, + gate_low_rank_dim: int = 128, + a_low_rank_dim: int = 64, + v_low_rank_dim: int = 16, + hidden_act: str = "sqrelu", + max_position_embeddings: int = 2048, + norm_first: bool = True, + norm_bias: bool = True, + norm_eps: float = 1e-5, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = True, + vocab_size: int = 32000, + value_dim: int | list[int] | None = None, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.norm_first = norm_first + self.num_hidden_layers = num_hidden_layers + + if head_dim is None and num_heads is not None: + head_dim = int(hidden_size // num_heads) + elif head_dim is not None and num_heads is None: + num_heads = int(hidden_size // head_dim) + + if value_dim is None: + value_dim = [hidden_size] * num_hidden_layers + elif isinstance(value_dim, int): + assert value_dim >= hidden_size, "value_dim must be greater than hidden_size" + assert value_dim % hidden_size == 0, "value_dim must be divisible by hidden_size" + value_dim = [value_dim] * num_hidden_layers + else: + assert len(value_dim) == num_hidden_layers, "value_dim must have the same length as num_hidden_layers" + for v in value_dim: + assert v >= hidden_size, "value_dim must be greater than hidden_size" + assert v % hidden_size == 0, "value_dim must be divisible by hidden_size" + + self.head_dim = head_dim + self.num_heads = num_heads + self.value_dim = value_dim + + self.decay_low_rank_dim = decay_low_rank_dim + self.gate_low_rank_dim = gate_low_rank_dim + self.a_low_rank_dim = a_low_rank_dim + self.v_low_rank_dim = v_low_rank_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.norm_bias = norm_bias + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + self.fuse_norm = fuse_norm + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/rwkv7/modeling_rwkv7.py b/fla/models/rwkv7/modeling_rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..d86b277f13b601edd7c54a6793d3aa9bd2da564b --- /dev/null +++ b/fla/models/rwkv7/modeling_rwkv7.py @@ -0,0 +1,545 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.rwkv7 import RWKV7Attention +from fla.models.rwkv7.configuration_rwkv7 import RWKV7Config +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, LayerNorm +from fla.modules.activations import ACT2FN +from fla.modules.l2warp import l2_warp +from fla.modules.token_shift import token_shift + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class RWKV7FeedForward(nn.Module): + + def __init__( + self, + hidden_size: int, + hidden_ratio: int | None = None, + intermediate_size: int | None = None, + hidden_act: str = 'sqrelu', + layer_idx: int = None, + num_hidden_layers: int = None, + ) -> RWKV7FeedForward: + super().__init__() + + self.hidden_size = hidden_size + if hidden_ratio is None: + hidden_ratio = 4 + if intermediate_size is None: + intermediate_size = int(hidden_size * hidden_ratio) + intermediate_size = 32 * ((intermediate_size + 32 - 1) // 32) + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + + self.x_k = nn.Parameter(torch.zeros(hidden_size)) + + self.key = nn.Linear(hidden_size, intermediate_size, bias=False) + self.value = nn.Linear(intermediate_size, hidden_size, bias=False) + self.act_fn = ACT2FN[hidden_act] + + self.layer_idx = layer_idx + self.num_hidden_layers = num_hidden_layers + + try: + from transformers.modeling_utils import _init_weights + except ImportError: + _init_weights = True + if _init_weights: + self.apply(self._initialize_weights) + for name, module in self.named_modules(): + module._in_rwkv_module = True + + def _initialize_weights(self, module: nn.Module): + if isinstance(module, RWKV7FeedForward): + with torch.no_grad(): + ratio_1_to_almost0 = 1.0 - (module.layer_idx / module.num_hidden_layers) # 1 to ~0 + ddd = torch.ones(1, 1, module.hidden_size) + for i in range(module.hidden_size): + ddd[0, 0, i] = i / module.hidden_size + module.x_k.data = 1.0 - torch.pow(ddd, ratio_1_to_almost0**4).squeeze() + + # Initialize key and value weights as in CMix_x070 + original_dtype = module.key.weight.dtype + module.key.weight.data = nn.init.orthogonal_(module.key.weight.data.to(torch.float32)).to(original_dtype) + module.value.weight.data.zero_() + + def forward( + self, + x: torch.Tensor, + attention_mask: torch.Tensor | None = None, + state: Cache | None = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> torch.Tensor: + if attention_mask is not None: + x = x.mul(attention_mask[:, -x.shape[-2]:, None]) + if state is not None: + delta, ffn_state = token_shift(x, cu_seqlens, cache=state[self.layer_idx]['ffn_state'], output_cache=True) + else: + delta, ffn_state = token_shift(x, cu_seqlens, output_cache=True) + if state is not None: + # no need to update the offset twice + state.update(ffn_state=ffn_state, layer_idx=self.layer_idx, offset=0) + return self.value(self.act_fn(self.key(x.addcmul(delta, self.x_k)))), state + + +class RWKV7Block(GradientCheckpointingLayer): + + def __init__( + self, + config: RWKV7Config, + layer_idx: int, + ) -> RWKV7Block: + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + if config.norm_first and layer_idx == 0: + self.pre_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + self.attn_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.attn = RWKV7Attention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_heads=config.num_heads, + decay_low_rank_dim=config.decay_low_rank_dim, + gate_low_rank_dim=config.gate_low_rank_dim, + a_low_rank_dim=config.a_low_rank_dim, + v_low_rank_dim=config.v_low_rank_dim, + norm_eps=config.norm_eps, + fuse_norm=config.fuse_norm, + layer_idx=layer_idx, + value_dim=config.value_dim[layer_idx], + num_hidden_layers=config.num_hidden_layers, + ) + self.ffn_norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + self.ffn = RWKV7FeedForward( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + layer_idx=layer_idx, + num_hidden_layers=config.num_hidden_layers, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + v_first: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = self.pre_norm(hidden_states) if hasattr(self, 'pre_norm') else hidden_states + hidden_states = self.attn_norm(residual) + hidden_states, attentions, past_key_values, v_first = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + v_first=v_first, + cu_seqlens=cu_seqlens, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.ffn_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.ffn_norm(hidden_states) + hidden_states, past_key_values = self.ffn( + hidden_states, attention_mask, past_key_values, cu_seqlens, **kwargs, + ) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values, v_first) + + return outputs + + +class RWKV7PreTrainedModel(PreTrainedModel): + + config_class = RWKV7Config + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['RWKV7Block'] + _supports_cache_class = True + _skip_keys_device_placement = ["past_key_values"] + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + @torch.no_grad() + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = True, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, nn.Embedding): + # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/train_temp/src/model.py#L396C12-L399C58 + scale = -1e-4 + nn.init.uniform_(module.weight, a=scale, b=-scale) + elif isinstance(module, nn.Linear) and hasattr(self, 'lm_head') and module is self.lm_head: + # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/train_temp/src/model.py#L403 + if self.config.vocab_size > self.config.hidden_size: + scale = 0.5 * math.sqrt(self.config.vocab_size / self.config.hidden_size) + else: + scale = 0.5 + original_dtype = module.weight.dtype + module.weight.data = nn.init.orthogonal_(module.weight.data.to(torch.float32), gain=scale).to(original_dtype) + # Init Attention parameters + elif isinstance(module, (nn.Linear, nn.Conv1d)) and getattr(module, '_in_rwkv_module', False) is False: + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Parameter): + nn.init.normal_(module, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters') and getattr(module, '_in_rwkv_module', False) is False: + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class RWKV7Model(RWKV7PreTrainedModel): + + def __init__(self, config: RWKV7Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([RWKV7Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (LayerNorm if config.fuse_norm else nn.LayerNorm)( + config.hidden_size, + bias=config.norm_bias, + eps=config.norm_eps, + ) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def load_state_dict(self, state_dict, strict=True, assign=False): + """ + Override the load_state_dict method to handle migration from version 1 to version 2. + Handles hierarchical keys like 'model.layers.0.attn.x_x'. + """ + # Collect all layer indices from the state_dict keys + layer_indices = set() + for key in state_dict.keys(): + if key.startswith("model.layers."): + # Extract the layer index from the key + try: + layer_idx = int(key.split(".")[2]) # Extract the number after 'model.layers.' + layer_indices.add(layer_idx) + except ValueError: + # Skip keys that don't match the expected format + continue + + # Sort the layer indices to process them in order + sorted_layer_indices = sorted(layer_indices) + + # Migration logic for each layer + for layer_idx in sorted_layer_indices: + layer_prefix = f"model.layers.{layer_idx}" + attn_prefix = f"{layer_prefix}.attn" + + # Check if the layer contains the old 'x_x' parameter + if f"{attn_prefix}.x_x" in state_dict: + logger.info(f"Migrating weights for layer {layer_idx} from RWKV7Attention version 1 to version 2...") + # Extract the x_x parameter + x_x = state_dict[f"{attn_prefix}.x_x"] + with torch.no_grad(): + # Create new parameters for version 2 + state_dict[f"{attn_prefix}.x_r"] = x_x[0].unsqueeze(0).unsqueeze(0) + state_dict[f"{attn_prefix}.x_w"] = x_x[1].unsqueeze(0).unsqueeze(0) + state_dict[f"{attn_prefix}.x_k"] = x_x[2].unsqueeze(0).unsqueeze(0) + state_dict[f"{attn_prefix}.x_v"] = x_x[3].unsqueeze(0).unsqueeze(0) + state_dict[f"{attn_prefix}.x_a"] = x_x[4].unsqueeze(0).unsqueeze(0) + state_dict[f"{attn_prefix}.x_g"] = x_x[5].unsqueeze(0).unsqueeze(0) + + # Call the parent method to load the modified state_dict + try: + super().load_state_dict(state_dict, strict=strict, assign=assign) + except TypeError: + # If the parent method does not support `assign`, fall back to strict loading + logger.warning( + "`assign` parameter is not supported by the parent `load_state_dict` method. " + "Falling back to default behavior.", + ) + super().load_state_dict(state_dict, strict=strict) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + cu_seqlens: torch.LongTensor | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`RWKV7Model` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + + v_first = torch.zeros_like(hidden_states) + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values, v_first = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + v_first=v_first, + cu_seqlens=cu_seqlens, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class RWKV7ForCausalLM(RWKV7PreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = RWKV7Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | None = None, + labels: torch.LongTensor | None = None, + shift_labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, logits = None, None + has_labels = (labels is not None) or (shift_labels is not None) + if not (self.config.fuse_linear_cross_entropy and has_labels): + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if has_labels: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + + # shift_labels: See https://github.com/huggingface/transformers/pull/36607/files. + if shift_labels is None: + shift_labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + shift_labels = shift_labels.to(hidden_states.device) + + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, shift_labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(shift_labels.numel(), -1), shift_labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/samba/__init__.py b/fla/models/samba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ce3f26751b5dbd4fc20c68970eaefd3b3a91c2df --- /dev/null +++ b/fla/models/samba/__init__.py @@ -0,0 +1,11 @@ +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.samba.configuration_samba import SambaConfig +from fla.models.samba.modeling_samba import SambaForCausalLM, SambaModel + +AutoConfig.register(SambaConfig.model_type, SambaConfig, exist_ok=True) +AutoModel.register(SambaConfig, SambaModel, exist_ok=True) +AutoModelForCausalLM.register(SambaConfig, SambaForCausalLM, exist_ok=True) + + +__all__ = ['SambaConfig', 'SambaForCausalLM', 'SambaModel'] diff --git a/fla/models/samba/configuration_samba.py b/fla/models/samba/configuration_samba.py new file mode 100644 index 0000000000000000000000000000000000000000..1ada5edb4680b9523664af9d17d7da38f9e61f85 --- /dev/null +++ b/fla/models/samba/configuration_samba.py @@ -0,0 +1,106 @@ + +import math +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class SambaConfig(PretrainedConfig): + + model_type = "samba" + + def __init__( + self, + hidden_size: int = 2304, + state_size: int = 16, + num_hidden_layers: int = 18, + norm_eps=1e-5, + pad_token_id: int = 0, + bos_token_id: int = 1, + eos_token_id: int = 2, + expand: int = 2, + conv_kernel: int = 4, + use_bias: bool = False, + use_conv_bias: bool = True, + hidden_act: str = "swish", + initializer_range: float = 0.02, + residual_in_fp32: bool = False, + time_step_rank: str = "auto", + time_step_scale: float = 1.0, + time_step_min: float = 0.001, + time_step_max: float = 0.1, + time_step_init_scheme: str = "random", + time_step_floor: float = 1e-4, + max_position_embeddings: int = 2048, + attn: dict | None = { + 'layers': (1, 3, 5, 7, 9, 11, 13, 15, 17), + 'num_heads': 18, + 'num_kv_heads': 18, + 'qkv_bias': False, + 'window_size': 2048, + 'rope_theta': 10000., + }, + hidden_ratio: int | None = 4, + rescale_prenorm_residual: bool = False, + use_cache: bool = True, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + tie_word_embeddings: bool = False, + **kwargs, + ): + self.hidden_size = hidden_size + self.state_size = state_size + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.conv_kernel = conv_kernel + self.expand = expand + self.intermediate_size = int(expand * self.hidden_size) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.use_bias = use_bias + self.use_conv_bias = use_conv_bias + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.time_step_rank = math.ceil(self.hidden_size / 16) if time_step_rank == "auto" else time_step_rank + self.time_step_scale = time_step_scale + self.time_step_min = time_step_min + self.time_step_max = time_step_max + self.time_step_init_scheme = time_step_init_scheme + self.time_step_floor = time_step_floor + self.max_position_embeddings = max_position_embeddings + self.attn = attn + self.hidden_ratio = hidden_ratio + self.rescale_prenorm_residual = rescale_prenorm_residual + self.residual_in_fp32 = residual_in_fp32 + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/samba/modeling_samba.py b/fla/models/samba/modeling_samba.py new file mode 100644 index 0000000000000000000000000000000000000000..4ca6d8938445f85491151f61bd9dbe35154be2c6 --- /dev/null +++ b/fla/models/samba/modeling_samba.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +from torch import nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.mamba import Mamba +from fla.models.samba.configuration_samba import SambaConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as SambaMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class SambaBlock(GradientCheckpointingLayer): + + def __init__(self, config, layer_idx): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.mixer_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + if config.attn is not None and layer_idx in config.attn['layers']: + self.mixer = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + else: + self.mixer = Mamba( + hidden_size=config.hidden_size, + state_size=config.state_size, + conv_kernel=config.conv_kernel, + intermediate_size=config.intermediate_size, + time_step_rank=config.time_step_rank, + use_bias=config.use_bias, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = SambaMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ): + residual = hidden_states + hidden_states = self.mixer_norm(hidden_states) + hidden_states, attentions, past_key_values = self.mixer( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + return hidden_states, attentions, past_key_values + + +class SambaPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = SambaConfig + base_model_prefix = "backbone" + _no_split_modules = ["SambaBlock"] + supports_gradient_checkpointing = True + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + if not getattr(module.bias, "_no_reinit", False): + nn.init.zeros_(module.bias) + elif isinstance(module, Mamba): + module.A_log._no_weight_decay = True + module.D._no_weight_decay = True + + dt_init_std = self.config.time_step_rank**-0.5 * self.config.time_step_scale + if self.config.time_step_init_scheme == "constant": + nn.init.constant_(module.dt_proj.weight, dt_init_std) + elif self.config.time_step_init_scheme == "random": + nn.init.uniform_(module.dt_proj.weight, -dt_init_std, dt_init_std) + + dt = torch.exp( + torch.rand(self.config.intermediate_size) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min), + ).clamp(min=self.config.time_step_floor) + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + module.dt_proj.bias.data = nn.Parameter(inv_dt.to(module.dt_proj.bias.device)) + module.dt_proj.bias._no_reinit = True + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name in ["out_proj.weight"]: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(self.config.num_layers) + + +class SambaModel(SambaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([SambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.gradient_checkpointing = False + self.norm_f = RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.LongTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one", + ) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for mixer_block in self.layers: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + hidden_states, attentions, past_key_values = mixer_block( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions and attentions is not None: + all_attns = all_attns + (attentions,) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns if all_attns else None, + ) + + +class SambaForCausalLM(SambaPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.backbone = SambaModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_input_embeddings(self): + return self.backbone.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.backbone.set_input_embeddings(new_embeddings) + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | CausalLMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.backbone( + input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + hidden_states = outputs[0] + + loss, logits = None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/transformer/__init__.py b/fla/models/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71b99803e4c654b02f1efec6a39b4731b8dc525a --- /dev/null +++ b/fla/models/transformer/__init__.py @@ -0,0 +1,12 @@ + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.transformer.configuration_transformer import TransformerConfig +from fla.models.transformer.modeling_transformer import TransformerForCausalLM, TransformerModel + +AutoConfig.register(TransformerConfig.model_type, TransformerConfig, exist_ok=True) +AutoModel.register(TransformerConfig, TransformerModel, exist_ok=True) +AutoModelForCausalLM.register(TransformerConfig, TransformerForCausalLM, exist_ok=True) + + +__all__ = ['TransformerConfig', 'TransformerForCausalLM', 'TransformerModel'] diff --git a/fla/models/transformer/configuration_transformer.py b/fla/models/transformer/configuration_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..2c58acf195d6a15ed31e8a3e0cdac5c8972cdb74 --- /dev/null +++ b/fla/models/transformer/configuration_transformer.py @@ -0,0 +1,85 @@ + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class TransformerConfig(PretrainedConfig): + + model_type = 'transformer' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 24, + num_heads: int = 32, + num_kv_heads: int | None = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: int | None = None, + rope_theta: float | None = 10000., + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + initializer_range: float = 0.02, + elementwise_affine: bool | None = True, + norm_eps: float = 1e-6, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + + self.initializer_range = initializer_range + self.elementwise_affine = elementwise_affine + self.norm_eps = norm_eps + self.use_cache = use_cache + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/transformer/modeling_transformer.py b/fla/models/transformer/modeling_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..dcfa802f84d2493ad3deda1fc067f7ffac6a8463 --- /dev/null +++ b/fla/models/transformer/modeling_transformer.py @@ -0,0 +1,355 @@ +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.models.transformer.configuration_transformer import TransformerConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as TransformerMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class TransformerBlock(GradientCheckpointingLayer): + + def __init__(self, config: TransformerConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + qkv_bias=config.qkv_bias, + qk_norm=config.qk_norm, + window_size=config.window_size, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = TransformerMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + **kwargs: Unpack[Any], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attentions,) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +class TransformerPreTrainedModel(PreTrainedModel): + + config_class = TransformerConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['TransformerBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + rescale_prenorm_residual: bool = False, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + + +class TransformerModel(TransformerPreTrainedModel): + + def __init__( + self, + config: TransformerConfig, + ) -> TransformerModel: + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([TransformerBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + if output_attentions: + warnings.warn( + "`TransformerModel` does not support output attention weights now, so `output_attentions` is set to `False`.", + ) + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + # embed positions + hidden_states = inputs_embeds + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + next_cache = None + + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_attns] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_attns, + ) + + +class TransformerForCausalLM(TransformerPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = TransformerModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[Any], + ) -> tuple | CausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + logits = None if self.config.fuse_linear_cross_entropy else self.lm_head(hidden_states[:, -logits_to_keep:]) + + loss = None + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + # Enable model parallelism + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/models/utils.py b/fla/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9c671814a6205da6ad8546048f6fcb58fe1530c7 --- /dev/null +++ b/fla/models/utils.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +import inspect +from typing import Any + +import torch +import transformers +from packaging import version +from transformers.cache_utils import Cache as HFCacheBase +from transformers.generation import GenerationMixin +from transformers.utils.deprecation import deprecate_kwarg + +_TF_VERSION = transformers.__version__ +_NEED_NEW = "4.53.3" +_IS_TRANSFORMERS_4_56_PLUS = version.parse(_TF_VERSION) >= version.parse("4.56.0") + +if version.parse(_TF_VERSION) > version.parse(_NEED_NEW): + from transformers.cache_utils import CacheLayerMixin +else: + CacheLayerMixin = object + + +class FLALayer(CacheLayerMixin): + is_compileable = True + is_sliding = False + + def __init__(self): + super().__init__() + self.state = None + self._seen_tokens = 0 + + def lazy_initialization(self, key_states: torch.Tensor): + self.state = None + + def update( + self, + *, + recurrent_state: torch.Tensor | tuple[torch.Tensor, ...] | None = None, + attn_state: tuple[torch.Tensor, ...] | None = None, + conv_state: Any | None = None, + ffn_state: Any | None = None, + offset: int = 1, + cache_kwargs: dict[str, Any] | None = None, + **_: Any, + ) -> dict[str, Any]: + if cache_kwargs is None: + cache_kwargs = {} + window_size = cache_kwargs.get("window_size") + + if attn_state is not None and not isinstance(attn_state, (tuple, list)): + raise ValueError("`attn_state` must be a tuple/list of tensors") + + if self.state is None: + self.state = { + "recurrent_state": None, + "attn_state": None, + "conv_state": None, + "ffn_state": None, + } + + if recurrent_state is not None: + self.state["recurrent_state"] = recurrent_state + + # Extract input_size from attn_state if available (before potential window truncation) + has_attn_state = attn_state and attn_state[0] is not None + input_size = attn_state[0].shape[1] if has_attn_state else 0 + + if has_attn_state: + if self.state["attn_state"] is None: + if window_size is not None and input_size > window_size: + attn_state = tuple(x[:, -window_size:].contiguous() for x in attn_state) + self.state["attn_state"] = tuple(attn_state) + else: + old = self.state["attn_state"] + if window_size is not None and old[0].shape[1] >= window_size: + new_tuple = [] + for old_x, new_x in zip(old, attn_state, strict=False): + rolled = old_x.roll(-input_size, dims=1) + tail = new_x[:, -window_size:] + rolled[:, -tail.shape[1]:] = tail + new_tuple.append(rolled) + self.state["attn_state"] = tuple(new_tuple) + else: + self.state["attn_state"] = tuple( + torch.cat([old_x, new_x], dim=1) for old_x, new_x in zip(old, attn_state, strict=False) + ) + + if conv_state is not None: + self.state["conv_state"] = conv_state + if ffn_state is not None: + self.state["ffn_state"] = ffn_state + + if not hasattr(self, 'device'): + self.device = 'cpu' + for state in (recurrent_state, attn_state, conv_state, ffn_state): + if state is not None: + if isinstance(state, torch.Tensor): + self.device = state.device + elif isinstance(state, (tuple, list)): + first_tensor = next((item for item in state if isinstance(item, torch.Tensor)), None) + if first_tensor is not None: + self.device = first_tensor.device + elif hasattr(state, 'device'): + self.device = state.device + else: + # For custom state objects (e.g., LogLinearAttentionState), + # try to find a tensor attribute to get the device. + for attr in vars(state).values(): + if isinstance(attr, torch.Tensor): + self.device = attr.device + break + break + + # Track seen tokens from attn_state if available, otherwise use offset + if has_attn_state: + # Use input_size captured before potential window truncation + self._seen_tokens += input_size + else: + # For layers without attn_state (e.g., rwkv7, gated_deltanet), use offset + self._seen_tokens += offset + + return self.state + + def get_seq_length(self, cache_position=None) -> int: + return self._seen_tokens + + def get_max_cache_shape(self) -> int: + return -1 + + def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]: + return 0, 0 + + def offload(self): + if self.state is None: + return + + def to_cpu(x): + return x.to("cpu", non_blocking=True) if isinstance(x, torch.Tensor) else x + for k in ("recurrent_state", "attn_state", "conv_state", "ffn_state"): + v = self.state.get(k, None) + if v is None: + continue + if isinstance(v, (tuple, list)): + self.state[k] = tuple(to_cpu(t) for t in v) + else: + self.state[k] = to_cpu(v) + + def prefetch(self): + if self.state is None: + return + + def to_dev(x): + return x.to(self.device, non_blocking=True) if isinstance(x, torch.Tensor) else x + for k in ("recurrent_state", "attn_state", "conv_state", "ffn_state"): + v = self.state.get(k, None) + if v is None: + continue + if isinstance(v, (tuple, list)): + self.state[k] = tuple(to_dev(t) for t in v) + else: + self.state[k] = to_dev(v) + + def reset(self): + pass + + +class LegacyFLACache(HFCacheBase): + """ + A cache used for storing hidden states produced by flash linear attention models. + + It stores the states of each layer as the tensor of shape `[batch_size, key_dim, value_dim]`. + """ + + is_compileable = True + + def __init__( + self, + seen_tokens: int = 0, + ) -> LegacyFLACache: + super().__init__() + + self.states: list[dict[str, Any]] = [] + + self._seen_tokens = seen_tokens # Used in `generate` to keep tally of how many tokens the cache has seen + + def __getitem__(self, layer_idx: int) -> dict[str, Any]: + if layer_idx < len(self): + return self.states[layer_idx] + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __iter__(self): + yield from self.states + + def __len__(self): + return len(self.states) + + def update( + self, + recurrent_state: tuple[torch.Tensor] | None = None, + attn_state: tuple[torch.Tensor] | None = None, + conv_state: tuple[torch.Tensor] | None = None, + ffn_state: tuple[torch.Tensor] | None = None, + layer_idx: int = 0, + offset: int | None = 1, + cache_kwargs: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """ + Args: + recurrent_state (`torch.Tensor`): + The new recurrent state to cache. + attn_state (`tuple[torch.Tensor]`): + The new attention key/value states to cache. + conv_state (`tuple[torch.Tensor]`): + The new convolution state to cache. + ffn_state (`tuple[torch.Tensor]`): + The new feed-forward state to cache. + layer_idx (`int`, defaults to 0): + The index of the layer to cache the states for. + offset (`int`, defaults to 1): + The number of new tokens being processed. + cache_kwargs (`Dict[str, Any]`): + Additional arguments for the cache subclass. + + Return: + Dictionary of the updated state. + """ + + if cache_kwargs is None: + cache_kwargs = {} + if attn_state is not None: + input_size = attn_state[0].shape[1] + window_size = cache_kwargs.get('window_size') + if not isinstance(attn_state, (tuple, list)): + raise ValueError("`attn_state` must be a tuple of tensors for key/value states") + if len(self.states) <= layer_idx: + # update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += offset + if attn_state is not None: + if window_size is not None and input_size > window_size: + attn_state = [state[:, -window_size:].contiguous() for state in attn_state] + state = dict( + recurrent_state=recurrent_state, + attn_state=attn_state, + conv_state=conv_state, + ffn_state=ffn_state, + ) + self.states.append(state) + else: + # update the number of seen tokens + if layer_idx == len(self.states) - 1: + self._seen_tokens += offset + state = self.states[layer_idx] + if recurrent_state is not None: + state['recurrent_state'] = recurrent_state + if attn_state is not None: + if window_size is not None and state['attn_state'][0].shape[1] == window_size: + for i, (old_state, new_state) in enumerate(zip(state['attn_state'], attn_state, strict=False)): + # DO NOT allocate new memory if the cache is full + # roll the key/value states to the left by `input_size` + old_state = old_state.roll(-input_size, 1) + # replace the last `input_size` tokens with the new key/value states + old_state[:, -input_size:] = new_state + state['attn_state'][i] = old_state + else: + attn_state = [ + torch.cat([old_state, new_state], 1) + for old_state, new_state in zip(state['attn_state'], attn_state, strict=False) + ] + state['attn_state'] = attn_state + if conv_state is not None: + state['conv_state'] = conv_state + if ffn_state is not None: + state['ffn_state'] = ffn_state + + return state + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + if len(self.states) <= layer_idx: + return 0 + return self._seen_tokens + + def get_max_cache_shape(self) -> int | None: + """Returns the maximum sequence length of the cached states. Cache does not have a maximum length.""" + return None + + def to_legacy_cache(self) -> tuple: + return tuple(self.states) + + @classmethod + @torch.compiler.disable + def from_legacy_cache( + cls, + past_key_values: tuple | None = None, + seen_tokens: int = 0, + ) -> LegacyFLACache: + """Converts a cache in the legacy cache format into an equivalent `Cache`.""" + + cache = cls(seen_tokens) + if isinstance(past_key_values, list): + for layer_idx in range(len(past_key_values)): + cache.states.append(past_key_values[layer_idx]) + return cache + + +class FLACache(HFCacheBase): + """ + A cache used for storing hidden states produced by flash linear attention models. + + It stores the states of each layer as the tensor of shape `[batch_size, key_dim, value_dim]`. + """ + + is_compileable = True + + def __init__(self, seen_tokens: int = 0, **kwargs): + parent_init = super().__init__ + sig = inspect.signature(parent_init) + param_names = list(sig.parameters.keys()) + + if 'layer_class_to_replicate' in param_names: + self.use_layer_class_to_replicate = True + super().__init__(layer_class_to_replicate=FLALayer, **kwargs) + elif 'layer_classes' in param_names: + self.use_layer_class_to_replicate = False + super().__init__(layer_classes=FLALayer, **kwargs) + else: + raise TypeError( + "FLA cache initialization failed: HFCacheBase.__init__ accepts neither " + "'layer_class_to_replicate' nor 'layer_classes'. This might be caused by an incompatible " + "transformers version. Please check your transformers>=4.36.0", + ) + self._seen_tokens = int(seen_tokens) + + def update( + self, + recurrent_state: tuple[torch.Tensor] | None = None, + attn_state: tuple[torch.Tensor] | None = None, + conv_state: tuple[torch.Tensor] | None = None, + ffn_state: tuple[torch.Tensor] | None = None, + layer_idx: int = 0, + offset: int | None = 1, + cache_kwargs: dict[str, Any] | None = None, + ) -> dict[str, Any]: + if not self.use_layer_class_to_replicate: + self.append_new_layers(layer_idx) + else: + while len(self.layers) <= layer_idx: + self.layers.append(self.layer_class_to_replicate()) + # Per-layer seen_tokens is now tracked in FLALayer.update() + + return self.layers[layer_idx].update( + recurrent_state=recurrent_state, + attn_state=attn_state, + conv_state=conv_state, + ffn_state=ffn_state, + offset=offset if offset is not None else 1, + cache_kwargs=cache_kwargs, + ) + + def __getitem__(self, layer_idx: int) -> dict[str, Any]: + if layer_idx >= len(self.layers): + raise KeyError(f"Cache only have {len(self.layers)} layers, however accessed {layer_idx} out of bounds") + return self.layers[layer_idx].state + + def __iter__(self): + for i in range(len(self.layers)): + yield self[i] + + def __len__(self): + return super().__len__() + + def get_seq_length(self, layer_idx: int | None = 0, cache_position=None) -> int: + if len(self.layers) <= (layer_idx or 0): + return 0 + return self.layers[layer_idx or 0].get_seq_length() + + def get_max_cache_shape(self, layer_idx: int = 0) -> int: + return -1 + + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + # kv_length = past_seen + current_query_length + query_len = int(cache_position.shape[0]) if cache_position is not None else 0 + kv_length = int(self.get_seq_length(layer_idx)) + query_len + return kv_length, 0 + + def to_legacy_cache(self) -> tuple[dict[str, Any], ...]: + return tuple(self[i] for i in range(len(self.layers))) + + @classmethod + @torch.compiler.disable + def from_legacy_cache( + cls, + past_key_values: tuple[dict[str, Any], ...] | None = None, + seen_tokens: int = 0, + **kwargs, + ) -> FLACache: + cache = cls(seen_tokens=seen_tokens, **kwargs) + if isinstance(past_key_values, (list, tuple)): + for i, st in enumerate(past_key_values): + while len(cache.layers) <= i: + cache.layers.append(cache.layer_class_to_replicate()) + cache.layers[i].state = dict(st) + return cache + + +class FLAGenerationMixin(GenerationMixin): + """ + Flash Linear Attention Generation Mixin that provides version-compatible generation methods. + This mixin handles transformers library version differences, particularly for prepare_inputs_for_generation. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor = None, + past_key_values: HFCacheBase | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool = True, + logits_to_keep: int | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs, + ): + # Use pre-computed version comparison for performance + if _IS_TRANSFORMERS_4_56_PLUS: + # For transformers 4.56.0+, use cache_position-based logic + model_inputs = {} + + # Handle cache-dependent input preparation + if past_key_values is not None: + model_inputs["past_key_values"] = past_key_values + + # Use the new cache-dependent input preparation method if available + if hasattr(self, '_cache_dependant_input_preparation') and cache_position is not None: + inputs_embeds, input_ids = self._cache_dependant_input_preparation( + input_ids, inputs_embeds, cache_position, + ) + elif cache_position is not None: + # Fallback: manually slice using cache_position + if input_ids is not None and input_ids.shape[1] != cache_position.shape[0]: + input_ids = input_ids[:, cache_position] + elif hasattr(past_key_values, '__len__') and len(past_key_values) > 0: + # Ultimate fallback to old behavior + input_ids = input_ids[:, -1:] + + # Handle input format (similar to base class logic) + if inputs_embeds is not None and (cache_position is None or len(cache_position) == inputs_embeds.shape[1]): + model_inputs['inputs_embeds'] = inputs_embeds + model_inputs['input_ids'] = None + else: + model_inputs['input_ids'] = input_ids.contiguous() if input_ids is not None else None + model_inputs['inputs_embeds'] = None + + model_inputs['cache_position'] = cache_position + + else: + # For older transformers versions, use the original logic + model_inputs = {} + # only last token for `inputs_ids` if the `past_key_values` is not empty. + if past_key_values is not None and hasattr(past_key_values, '__len__') and len(past_key_values) > 0: + input_ids = input_ids[:, -1:] + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and hasattr(past_key_values, '__len__') and len(past_key_values) == 0: + model_inputs = {'inputs_embeds': inputs_embeds} + else: + # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise + # recompiles graphs as the stride of the inputs is a guard. + # Ref: https://github.com/huggingface/transformers/pull/29114 + # TODO: use `next_tokens` directly instead. + model_inputs = {'input_ids': input_ids.contiguous()} + + if logits_to_keep is not None: + model_inputs['logits_to_keep'] = logits_to_keep + + model_inputs.update({ + 'past_key_values': past_key_values, + 'use_cache': use_cache, + 'attention_mask': attention_mask, + }) + return model_inputs + + +if version.parse(_TF_VERSION) > version.parse(_NEED_NEW): + class Cache(FLACache): + def __init__(self, seen_tokens: int = 0, **kwargs: Any) -> None: + super().__init__(seen_tokens=seen_tokens, **kwargs) +else: + class Cache(LegacyFLACache): + def __init__(self, seen_tokens: int = 0, **kwargs: Any) -> None: + super().__init__(seen_tokens=seen_tokens) diff --git a/fla/modules/__init__.py b/fla/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..030aa7f750181112ed4da5ede59f2ffbdd9bbdcd --- /dev/null +++ b/fla/modules/__init__.py @@ -0,0 +1,25 @@ +# Minimal exports for the Quasar/Bailing training path. +# Avoid eagerly importing optional model losses and long-convolution modules. + +from fla.modules.convolution import ShortConvolution +from fla.modules.fused_norm_gate import FusedRMSNormGated +from fla.modules.layernorm import RMSNorm, GroupNorm, LayerNorm +from fla.modules.rotary import RotaryEmbedding +from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss +from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss +from fla.modules.mlp import GatedMLP + +__all__ = [ + "FusedRMSNormGated", + "RMSNorm", + "GroupNorm", + "LayerNorm", + "RotaryEmbedding", + "ShortConvolution", + "FusedCrossEntropyLoss", + "FusedLinearCrossEntropyLoss", + "GatedMLP", +] + + + diff --git a/fla/modules/activations.py b/fla/modules/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..cae3f8a8a690752d2d2ce68291b2acb7d5fa3a72 --- /dev/null +++ b/fla/modules/activations.py @@ -0,0 +1,743 @@ +# Copyright (c) 2023-2025, Tri Dao, Yu Zhang, Songlin Yang. + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +def _get_stride(x: torch.Tensor) -> int: + """Get the row stride for viewing a tensor as 2D (num_rows, D) where D = shape[-1]. + + Returns stride(-2) if the tensor is at least 2D, or 0 for 1D tensors. + The caller must ensure the tensor is "inner-contiguous" (stride(-1) == 1 and + higher dims are contiguous relative to dim -2) before using this value. + """ + if x.ndim < 2: + return 0 + return x.stride(-2) + + +def _is_inner_contiguous(x: torch.Tensor) -> bool: + """Check if a tensor can be safely viewed as 2D (num_rows, D) with row stride = stride(-2). + + This holds when stride(-1) == 1 and all dimensions above -2 are contiguous + with respect to the dimension below them. + """ + ndim = x.ndim + if ndim < 2: + return True + if x.stride(-1) != 1: + return False + if ndim == 2: + # 2D: any layout with stride(-1)==1 is valid (can view as (T, D)) + return True + if ndim == 3: + # 3D (B, T, D): stride should be (T*D, D, 1) + return x.stride(0) == x.stride(-2) * x.shape[-2] + if ndim == 4: + # 4D (B, H, T, D): stride should be (H*T*D, T*D, D, 1) + if x.stride(1) != x.stride(-2) * x.shape[-2]: + return False + return x.stride(0) == x.stride(1) * x.shape[1] + # 5D+ fallback to loop + expected = x.stride(-2) * x.shape[-2] + for d in range(ndim - 3, -1, -1): + if x.stride(d) != expected: + return False + expected *= x.shape[d] + return True + + +def _ensure_inner_contiguous(x: torch.Tensor) -> torch.Tensor: + """Make the tensor inner-contiguous if it isn't already.""" + if _is_inner_contiguous(x): + return x + return x.contiguous() + + +def _alloc_output(x: torch.Tensor, contiguous: bool = False) -> torch.Tensor: + """Allocate output tensor: contiguous buffer or same layout as input.""" + if contiguous: + return x.new_empty(x.shape) + return torch.empty_like(x) + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def sigmoid_fwd_kernel( + x, y, + T, + D: tl.constexpr, + stride_x_row, + stride_y_row, + B: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + y_off = row * stride_y_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + y_val = 1.0 / (1.0 + exp(-x_val)) + tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def sigmoid_bwd_kernel( + x, dy, dx, + T, + D: tl.constexpr, + stride_x_row, + stride_dy_row, + stride_dx_row, + B: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + dy_off = row * stride_dy_row + col + dx_off = row * stride_dx_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32) + s = 1.0 / (1.0 + exp(-x_val)) + dx_val = g_val * s * (1.0 - s) + tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask) + + +@torch.compiler.disable +def sigmoid_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + T, D = x.numel(), x.shape[-1] + y = _alloc_output(x, output_contiguous) + sigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, y, T=T, D=D, + stride_x_row=_get_stride(x), + stride_y_row=_get_stride(y), + ) + return y + + +@torch.compiler.disable +def sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + dy = _ensure_inner_contiguous(dy) + T, D = x.numel(), x.shape[-1] + dx = _alloc_output(x, output_contiguous) + sigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, dy, dx, T=T, D=D, + stride_x_row=_get_stride(x), + stride_dy_row=_get_stride(dy), + stride_dx_row=_get_stride(dx), + ) + return dx + + +class SigmoidFunction(torch.autograd.Function): + + @staticmethod + @input_guard(no_guard_contiguous=True) + def forward(ctx, x): + ctx.save_for_backward(x) + return sigmoid_fwd(x) + + @staticmethod + @input_guard(no_guard_contiguous=True) + def backward(ctx, dout): + x, = ctx.saved_tensors + return sigmoid_bwd(x, dout) + + +sigmoid = SigmoidFunction.apply + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def logsigmoid_fwd_kernel( + x, + y, + temperature, + T, + D: tl.constexpr, + stride_x_row, + stride_y_row, + B: tl.constexpr, +): + i = tl.program_id(0) + o_i = i * B + tl.arange(0, B) + m_i = o_i < T + row = o_i // D + col = o_i % D + x_off = row * stride_x_row + col + y_off = row * stride_y_row + col + + b_x = tl.load(x + x_off, mask=m_i, other=0.).to(tl.float32) + b_m = tl.minimum(0., b_x) + b_z = 1. + exp(-tl.abs(b_x)) + b_y = (b_m - log(b_z)) / temperature + tl.store(y + y_off, b_y.to(y.dtype.element_ty), mask=m_i) + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def logsigmoid_bwd_kernel( + x, + dx, + dy, + temperature, + T, + D: tl.constexpr, + stride_x_row, + stride_dx_row, + stride_dy_row, + B: tl.constexpr, +): + i = tl.program_id(0) + o_i = i * B + tl.arange(0, B) + m_i = o_i < T + row = o_i // D + col = o_i % D + x_off = row * stride_x_row + col + dx_off = row * stride_dx_row + col + dy_off = row * stride_dy_row + col + + b_x = tl.load(x + x_off, mask=m_i, other=0.).to(tl.float32) + b_dy = tl.load(dy + dy_off, mask=m_i, other=0.).to(tl.float32) + b_dx = b_dy * ((1. - tl.sigmoid(b_x)) / temperature) + tl.store(dx + dx_off, b_dx.to(dx.dtype.element_ty), mask=m_i) + + +@torch.compiler.disable +def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + T, D = x.numel(), x.shape[-1] + y = _alloc_output(x, output_contiguous) + logsigmoid_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x=x, + y=y, + temperature=temperature, + T=T, + D=D, + stride_x_row=_get_stride(x), + stride_y_row=_get_stride(y), + ) + return y + + +@torch.compiler.disable +def logsigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + dy = _ensure_inner_contiguous(dy) + T, D = x.numel(), x.shape[-1] + dx = _alloc_output(x, output_contiguous) + logsigmoid_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x=x, + dx=dx, + dy=dy, + temperature=temperature, + T=T, + D=D, + stride_x_row=_get_stride(x), + stride_dx_row=_get_stride(dx), + stride_dy_row=_get_stride(dy), + ) + return dx + + +class LogSigmoidFunction(torch.autograd.Function): + + @staticmethod + @input_guard(no_guard_contiguous=True) + def forward(ctx, x, temperature): + ctx.save_for_backward(x) + ctx.temperature = temperature + return logsigmoid_fwd(x, temperature) + + @staticmethod + @input_guard(no_guard_contiguous=True) + def backward(ctx, dy): + x, = ctx.saved_tensors + return logsigmoid_bwd(x, dy, ctx.temperature), None + + +def logsigmoid(x: torch.Tensor, temperature: float = 1.) -> torch.Tensor: + return LogSigmoidFunction.apply(x, temperature) + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def swish_fwd_kernel( + x, y, + T, + D: tl.constexpr, + stride_x_row, + stride_y_row, + B: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + y_off = row * stride_y_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + s = 1.0 / (1.0 + exp(-x_val)) + y_val = x_val * s + tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def swish_bwd_kernel( + x, dy, dx, + T, + D: tl.constexpr, + stride_x_row, + stride_dy_row, + stride_dx_row, + B: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + dy_off = row * stride_dy_row + col + dx_off = row * stride_dx_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32) + s = 1.0 / (1.0 + exp(-x_val)) + dx_val = g_val * s * (1.0 + x_val * (1.0 - s)) + tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask) + + +@torch.compiler.disable +def swish_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + T, D = x.numel(), x.shape[-1] + y = _alloc_output(x, output_contiguous) + swish_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, y, T=T, D=D, + stride_x_row=_get_stride(x), + stride_y_row=_get_stride(y), + ) + return y + + +@torch.compiler.disable +def swish_bwd(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: + x = _ensure_inner_contiguous(x) + dy = _ensure_inner_contiguous(dy) + T, D = x.numel(), x.shape[-1] + dx = _alloc_output(x, output_contiguous) + swish_bwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, dy, dx, T=T, D=D, + stride_x_row=_get_stride(x), + stride_dy_row=_get_stride(dy), + stride_dx_row=_get_stride(dx), + ) + return dx + + +class SwishFunction(torch.autograd.Function): + + @staticmethod + @input_guard(no_guard_contiguous=True) + def forward(ctx, x): + ctx.save_for_backward(x) + return swish_fwd(x) + + @staticmethod + @input_guard(no_guard_contiguous=True) + def backward(ctx, dout): + x, = ctx.saved_tensors + return swish_bwd(x, dout) + + +swish = SwishFunction.apply + +# 1/sqrt(2*pi)-> 0.3989423 +# 1/sqrt(2) -> 0.70710678 +# sqrt(2/pi) -> 0.79788456 + + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.compile +def bias_gelu(y, bias): + x = bias + y + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.compile +def bias_gelu_bwd(g, y, bias): + """Assume that y has shape (B, D=D) and bias has shape (D)""" + x = bias + y + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + grad_y = ff * g + return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) + + +class GeLUFunction(torch.autograd.Function): + + @staticmethod + # bias is an optional argument + def forward(ctx, input, bias): + ctx.save_for_backward(input, bias) + return bias_gelu(input, bias) + + @staticmethod + def backward(ctx, grad_output): + input, bias = ctx.saved_tensors + tmp = bias_gelu_bwd(grad_output, input, bias) + return tmp, tmp + + +bias_gelu_impl = GeLUFunction.apply + + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.compile +def gelu_fwd(x): + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.compile +def gelu_bwd(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + return (ff * g).to(dtype=x.dtype) + + +class FastGeLUFunction(torch.autograd.Function): + @staticmethod + # bias is an optional argument + def forward(ctx, input): + ctx.save_for_backward(input) + return gelu_fwd(input) + + @staticmethod + def backward(ctx, grad_output): + (input,) = ctx.saved_tensors + tmp = gelu_bwd(grad_output, input) + return tmp + + +fast_gelu_impl = FastGeLUFunction.apply + + +@torch.compile +def relu_bwd(g, x): + return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) + + +@torch.compile +def sqrelu_fwd(x): + r = F.relu(x.float()) + return (r * r).to(dtype=x.dtype) + + +@torch.compile +def sqrelu_bwd(g, x): + return (2.0 * g * F.relu(x.float())).to(dtype=x.dtype) + + +class SquaredReLUFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, input): + ctx.save_for_backward(input) + return sqrelu_fwd(input) + + @staticmethod + def backward(ctx, grad_output): + input, = ctx.saved_tensors + return sqrelu_bwd(grad_output, input) + + +sqrelu = SquaredReLUFunction.apply + + +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def swiglu_fwd_kernel( + x, y, z, + T, + D: tl.constexpr, + stride_x_row, + stride_y_row, + stride_z_row, + B: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + y_off = row * stride_y_row + col + z_off = row * stride_z_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + y_val = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32) + s = 1.0 / (1.0 + exp(-x_val)) + z_val = x_val * s * y_val + tl.store(z + z_off, z_val.to(z.dtype.element_ty), mask=mask) + + +@triton.heuristics({ + 'HAS_WEIGHT': lambda args: args['z'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'B': bs}, num_warps=num_warps) + for bs in [512, 1024, 2048, 4096, 8192] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def swiglu_fwdbwd_kernel( + x, y, g, dx, dy, z, + T, + D: tl.constexpr, + stride_x_row, + stride_y_row, + stride_g_row, + stride_dx_row, + stride_dy_row, + stride_z_row, + B: tl.constexpr, + HAS_WEIGHT: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * B + tl.arange(0, B) + mask = offs < T + row = offs // D + col = offs % D + x_off = row * stride_x_row + col + y_off = row * stride_y_row + col + g_off = row * stride_g_row + col + dx_off = row * stride_dx_row + col + dy_off = row * stride_dy_row + col + x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32) + y_val = tl.load(y + y_off, mask=mask, other=0.).to(tl.float32) + g_val = tl.load(g + g_off, mask=mask, other=0.).to(tl.float32) + + s = 1.0 / (1.0 + exp(-x_val)) + x_s = x_val * s + dx_val = g_val * s * (1.0 + x_val * (1.0 - s)) * y_val + dy_val = g_val * x_s + + tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask) + tl.store(dy + dy_off, dy_val.to(dy.dtype.element_ty), mask=mask) + if HAS_WEIGHT: + z_off = row * stride_z_row + col + z_val = x_s * y_val + tl.store(z + z_off, z_val.to(z.dtype.element_ty), mask=mask) + + +@torch.compiler.disable +def swiglu_fwd(x: torch.Tensor, y: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: + assert x.shape == y.shape, f"swiglu_fwd: shape mismatch x={x.shape} y={y.shape}" + x = _ensure_inner_contiguous(x) + y = _ensure_inner_contiguous(y) + T, D = x.numel(), x.shape[-1] + z = _alloc_output(x, output_contiguous) + swiglu_fwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, y, z, T=T, D=D, + stride_x_row=_get_stride(x), + stride_y_row=_get_stride(y), + stride_z_row=_get_stride(z), + ) + return z + + +@torch.compiler.disable +def swiglu_fwdbwd( + x: torch.Tensor, + y: torch.Tensor, + g: torch.Tensor, + use_weight: bool = False, + output_contiguous: bool = False, +): + assert x.shape == y.shape == g.shape, f"swiglu_fwdbwd: shape mismatch x={x.shape} y={y.shape} g={g.shape}" + x = _ensure_inner_contiguous(x) + y = _ensure_inner_contiguous(y) + g = _ensure_inner_contiguous(g) + T, D = x.numel(), x.shape[-1] + dx = _alloc_output(x, output_contiguous) + dy = _alloc_output(y, output_contiguous) + if use_weight: + z = _alloc_output(x, output_contiguous) + else: + z = None + swiglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( + x, y, g, dx, dy, z, T=T, D=D, + stride_x_row=_get_stride(x), + stride_y_row=_get_stride(y), + stride_g_row=_get_stride(g), + stride_dx_row=_get_stride(dx), + stride_dy_row=_get_stride(dy), + stride_z_row=_get_stride(z) if z is not None else 0, + ) + if use_weight: + return dx, dy, z + return dx, dy + + +class SwiGLUFunction(torch.autograd.Function): + r""" + Swish-Gated Linear Unit (SwiGLU) function. + + .. math:: + \text{SwiGLU}(x, y) = swish(x) * y = \frac{x}{1 + \exp(-x)} * y + """ + + @staticmethod + @input_guard(no_guard_contiguous=True) + def forward(ctx, x, y): + ctx.save_for_backward(x, y) + return swiglu_fwd(x, y) + + @staticmethod + @input_guard(no_guard_contiguous=True) + def backward(ctx, dout): + x, y = ctx.saved_tensors + return swiglu_fwdbwd(x, y, dout) + + +class SwiGLULinearFunction(torch.autograd.Function): + r""" + Swish-Gated Linear Unit (SwiGLU) function followed by a linear transformation. + + .. math:: + \text{SwiGLULinear}(x, y, W, b) = (swish(x) * y) W + b + + This simple wrap discards the intermediate results of SwiGLU(x, y) to save memory. + """ + + @staticmethod + @input_guard(no_guard_contiguous=True) + @autocast_custom_fwd + def forward(ctx, x, y, weight, bias): + z = swiglu_fwd(x, y, output_contiguous=True) + out = F.linear(z, weight, bias) + ctx.save_for_backward(x, y, weight) + ctx.linear_bias_is_none = bias is None + return out + + @staticmethod + @input_guard(no_guard_contiguous=True) + @autocast_custom_bwd + def backward(ctx, dout, *args): + x, y, weight = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dz = F.linear(dout, weight.t()).view_as(x) + dx, dy, z = swiglu_fwdbwd(x, y, dz, use_weight=True, output_contiguous=True) + dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1])) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + return dx, dy, dlinear_weight, dlinear_bias + + +swiglu = SwiGLUFunction.apply + + +swiglu_linear = SwiGLULinearFunction.apply + + +ACT2FN = { + 'relu': F.relu, + 'sigmoid': sigmoid, + 'logsigmoid': logsigmoid, + 'silu': swish, + 'swish': swish, + 'sqrelu': sqrelu, + 'gelu': fast_gelu_impl, + 'bias_gelu': bias_gelu_impl, +} diff --git a/fla/modules/conv/__init__.py b/fla/modules/conv/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f743ea465e8263760a41ec5a4ddbe6c4e1c305c --- /dev/null +++ b/fla/modules/conv/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from .causal_conv1d import causal_conv1d +from .long_conv import ImplicitLongConvolution, LongConvolution, PositionalEmbedding, fft_conv +from .short_conv import ShortConvolution + +__all__ = [ + 'ImplicitLongConvolution', + 'LongConvolution', + 'PositionalEmbedding', + 'ShortConvolution', + 'causal_conv1d', + 'fft_conv', +] diff --git a/fla/modules/conv/causal_conv1d.py b/fla/modules/conv/causal_conv1d.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae98b10a3996634ab1288ecd5046b42fae39e9b --- /dev/null +++ b/fla/modules/conv/causal_conv1d.py @@ -0,0 +1,124 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +"""Main interface for causal 1D convolution operations.""" + +import torch + +from fla.ops.cp import FLACPContext +from fla.utils import input_guard + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d( + x: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, + activation: str | None = None, + backend: str | None = 'triton', + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + """ + A causal 1D convolution implementation that powers Mamba/Mamba2 and DeltaNet architectures. + + When a residual connection is provided, this implements the Canon operation + described in the paper at https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5240330. + + Args: + x (torch.Tensor): + Input tensor of shape [B, T, D]. + weight (Optional[torch.Tensor]): + Weight tensor of shape [D, W]. Default: `None`. + bias (Optional[torch.Tensor]): + Bias tensor of shape [D]. Default: `None`. + residual (Optional[torch.Tensor]): + Residual tensor of shape [B, T, D]. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state tensor of shape [N, D, W], + where `N` is the number of sequences in the batch and `W` is the kernel size. + If provided, the initial state is used to initialize the cache. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape [N, D, W]. Default: `False`. + activation (Optional[str]): + Activations applied to output, only `swish`/`silu` or `None` (i.e., no activation) are supported. + Default: `None`. + backend (Optional[str]): + Specifies the backend to use for the convolution operation. Supported values are `'cuda'` 、 `'triton'` and `'mix'`. + Default: `'triton'`. + cu_seqlens (Optional[torch.Tensor]): + Cumulative sequence lengths (optional) + chunk_indices (Optional[torch.LongTensor]): + Chunk indices for variable-length sequences (optional) + + Returns: + Tuple of (output, final_state). + If `output_final_state` is `False`, the final state is `None`. + """ + # Import here to avoid circular dependencies + from fla.modules.conv.cp import causal_conv1d_cp + from fla.modules.conv.cuda import causal_conv1d_cuda, fast_causal_conv1d_fn + from fla.modules.conv.triton import CausalConv1dFunction + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + output = causal_conv1d_cp( + x=x, + weight=weight, + bias=bias, + activation=activation, + chunk_indices=chunk_indices, + cp_context=cp_context, + ) + return output, None + + if backend == 'triton': + y, final_state = CausalConv1dFunction.apply( + x, + weight, + bias, + residual, + initial_state, + output_final_state, + activation, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + ) + return y, final_state + elif backend == 'mix': + seq_idx = kwargs.get('seq_idx') + return fast_causal_conv1d_fn( + x, + weight, + bias, + residual, + initial_state, + output_final_state, + activation, + cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + seq_idx=seq_idx, + ) + elif backend == 'cuda': + return causal_conv1d_cuda( + x, + weight, + bias, + residual, + initial_state, + output_final_state, + activation, + cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + **kwargs, + ) + else: + raise ValueError(f"Unsupported backend: {backend}") diff --git a/fla/modules/conv/cp/__init__.py b/fla/modules/conv/cp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e2431fdc6f88b31f4778a29254cca25be491939d --- /dev/null +++ b/fla/modules/conv/cp/__init__.py @@ -0,0 +1,6 @@ +from .ops import CausalConv1dFunctionCP, causal_conv1d_cp + +__all__ = [ + 'CausalConv1dFunctionCP', + 'causal_conv1d_cp', +] diff --git a/fla/modules/conv/cp/ops.py b/fla/modules/conv/cp/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..06056c9a2e90a3792eed384c8ce8f67ba38b3cfe --- /dev/null +++ b/fla/modules/conv/cp/ops.py @@ -0,0 +1,251 @@ +import torch +import torch.distributed as dist + +from fla.ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd +from fla.ops.utils import prepare_chunk_indices + + +class CausalConv1dFunctionCP(torch.autograd.Function): + """ + Context Parallel version of CausalConv1dFunction. + + Forward: + 1. Get tails from previous rank to construct initial_state + 2. Call causal_conv1d_fwd + + Backward: + 1. Call causal_conv1d_bwd to get dx + 2. Sync communication: add next rank's first W-1 token gradients to current rank's last W-1 tokens + """ + + @staticmethod + def _prepare_initial_state_for_cp( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor | None, + context: FLACPContext, + group: dist.ProcessGroup | None, + ) -> torch.Tensor | None: + """Prepare initial_state for CP forward pass by communicating with previous rank. + + Args: + x: Input tensor of shape [1, T, D] + weight: Weight tensor of shape [D, W] + cu_seqlens: Cumulative sequence lengths + context: CP context + group: Process group for communication + + Returns: + initial_state: Initial state tensor of shape [N, D, W] or None + """ + if group is None: + return None + + W = weight.shape[-1] # weight: [D, W] + D = weight.shape[0] + initial_state = None + if not context.is_first_rank: + # Non-first rank needs initial_state + assert x.dim() == 3 and x.shape[0] == 1, f"CP requires [1, T, D], got {x.shape}" + x_2d = x.squeeze(0) # [T, D] + tails = x_2d[-(W-1):].contiguous() # [W-1, D] + heads = conv_cp_send_recv_fwd(tails, group) # [W-1, D] + # Construct initial_state: [N, D, W] + N = len(cu_seqlens) - 1 + initial_state = torch.zeros(N, D, W, device=x.device, dtype=x.dtype) + valid_len = min(W - 1, context.pre_num_conv_tokens) + if valid_len > 0: + # heads[-valid_len:]: [valid_len, D] -> [D, valid_len] + initial_state[0, :, -valid_len:] = heads[-valid_len:].T + else: + # First rank also needs to participate in communication (send tails) + x_2d = x.squeeze(0) + tails = x_2d[-(W-1):].contiguous() + _ = conv_cp_send_recv_fwd(tails, group) # Send but don't use + + return initial_state + + @staticmethod + def _correct_dx_for_cp( + dx: torch.Tensor, + dh0: torch.Tensor | None, + W: int, + group: dist.ProcessGroup | None, + is_first_rank: bool, + pre_num_conv_tokens: int = 0, + ) -> None: + """Correct dx gradients for CP backward pass by communicating with next rank. + + Args: + dx: Gradient tensor to be corrected, shape [1, T, D] + dh0: Gradient w.r.t. initial_state, shape [N, D, W] or None + W: Kernel size + group: Process group for communication + is_first_rank: Whether this is the first rank in the sequence's processing chain + pre_num_conv_tokens: Number of tokens from the previous rank that + belong to the first sequence on the current rank. Must match the + value used in the forward pass to construct initial_state. + """ + if group is None: + return + + D = dx.shape[-1] + # dh0: [N, D, W] or None + # We only care about the first sequence's initial_state gradient + if dh0 is not None: + # Only keep gradients for positions that had real data from the + # previous rank. The forward fills only the last valid_len positions + # of initial_state; gradients for the remaining (zero-padded) positions + # must not flow back, otherwise they leak into unrelated sequences. + valid_len = min(W - 1, pre_num_conv_tokens) + d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype) + if valid_len > 0: + d_initial_state[-valid_len:] = dh0[0, :, -valid_len:].T + else: + # dh0 is None only when this is the first rank (no initial_state needed) + assert is_first_rank, "dh0 should not be None when is_first_rank=False" + d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype) + # Sync communication: send d_initial_state to previous rank, receive from next rank + recv_d_init = conv_cp_send_recv_bwd(d_initial_state, group) # [W-1, D] + # Add to current rank's last W-1 tokens (these tokens are used as initial_state by next rank) + dx[0, -(W-1):, :].add_(recv_d_init) + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + activation: str | None, + chunk_indices: torch.Tensor | None, + cp_context: FLACPContext | None, + chunk_size: int | None, + backend: str = 'triton', + ): + # Import here to avoid circular dependency + from fla.modules.conv.triton.ops import causal_conv1d_fwd + + if cp_context is None: + raise ValueError("cp_context must be provided for CausalConv1dFunctionCP") + cu_seqlens = cp_context.cu_seqlens + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + group = cp_context.group + + # Get kernel_size + W = weight.shape[-1] # weight: [D, W] + # Prepare initial_state for CP + initial_state = CausalConv1dFunctionCP._prepare_initial_state_for_cp( + x=x, + weight=weight, + cu_seqlens=cu_seqlens, + context=cp_context, + group=group, + ) + + ctx.save_for_backward(x, weight, bias, initial_state) + ctx.activation = activation + ctx.cu_seqlens = cu_seqlens + ctx.cu_seqlens_cpu = cu_seqlens_cpu + ctx.chunk_indices = chunk_indices + ctx.chunk_size = chunk_size + ctx.group = group + ctx.W = W + ctx.is_first_rank = cp_context.is_first_rank + ctx.pre_num_conv_tokens = cp_context.pre_num_conv_tokens + + # Call original forward + y, _ = causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + residual=None, + initial_state=initial_state, + output_final_state=False, + activation=activation, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + BT=chunk_size, + ) + + return y + + @staticmethod + def backward(ctx, dy: torch.Tensor): + # Import here to avoid circular dependency + from fla.modules.conv.triton.ops import causal_conv1d_bwd + + x, weight, bias, initial_state = ctx.saved_tensors + group = ctx.group + W = ctx.W + + # Call original backward + dx, dw, db, _, dh0 = causal_conv1d_bwd( + x=x, + dy=dy, + dht=None, + weight=weight, + bias=bias, + residual=None, + initial_state=initial_state, + activation=ctx.activation, + cu_seqlens=ctx.cu_seqlens, + cu_seqlens_cpu=ctx.cu_seqlens_cpu, + chunk_indices=ctx.chunk_indices, + BT=ctx.chunk_size, + ) + + # Correct dx gradients for CP + CausalConv1dFunctionCP._correct_dx_for_cp( + dx=dx, + dh0=dh0, + W=W, + group=group, + is_first_rank=ctx.is_first_rank, + pre_num_conv_tokens=ctx.pre_num_conv_tokens, + ) + + return dx, dw, db, None, None, None, None, None + + +def causal_conv1d_cp( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + chunk_indices: torch.Tensor | None = None, + cp_context: FLACPContext | None = None, + chunk_size: int | None = None, + backend: str = 'triton', +): + """ + Context Parallel version of causal_conv1d. + + Automatically handles communication in CP environment: + - Forward: get initial_state from previous rank + - Backward: correct dx gradients + + Args: + x: Input tensor of shape [1, T, D] + weight: Weight tensor of shape [D, W] + bias: Bias tensor of shape [D] or None + activation: Activation function name or None + cu_seqlens: Cumulative sequence lengths + cu_seqlens_cpu: Cumulative sequence lengths on CPU + chunk_indices: Chunk indices for variable-length sequences + cp_context: CP context (required for CP mode) + """ + if cp_context is None: + raise ValueError("cp_context must be provided for causal_conv1d_cp") + + assert cp_context.conv1d_kernel_size is not None, "conv1d_kernel_size must be provided for causal_conv1d_cp" + assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for causal_conv1d_cp" + assert backend in ['triton'], "backend must be 'triton'" + chunk_size = chunk_size or 64 + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cp_context.cu_seqlens, chunk_size, cu_seqlens_cpu=cp_context.cu_seqlens_cpu) + + return CausalConv1dFunctionCP.apply( + x, weight, bias, activation, + chunk_indices, cp_context, chunk_size, backend + ) diff --git a/fla/modules/conv/cuda/__init__.py b/fla/modules/conv/cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7cc64781b0a721d20505f5013867cd91cc0d61e6 --- /dev/null +++ b/fla/modules/conv/cuda/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from .ops import FastCausalConv1dFn, causal_conv1d_cuda, fast_causal_conv1d_fn + +__all__ = [ + 'FastCausalConv1dFn', + 'causal_conv1d_cuda', + 'fast_causal_conv1d_fn', +] diff --git a/fla/modules/conv/cuda/ops.py b/fla/modules/conv/cuda/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..148f204a03cc81c5b6680f47f538971bcae49b18 --- /dev/null +++ b/fla/modules/conv/cuda/ops.py @@ -0,0 +1,228 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +"""CUDA-based mixed-mode implementation for causal convolution.""" + +import torch +from einops import rearrange + +from fla.modules.conv.triton import causal_conv1d_update_states +from fla.ops.utils import prepare_sequence_ids +from fla.utils import input_guard + +try: + from causal_conv1d.cpp_functions import causal_conv1d_bwd_function +except ImportError: + causal_conv1d_bwd_function = None + +try: + from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda +except ImportError: + causal_conv1d_fn_cuda = None + + +class FastCausalConv1dFn(torch.autograd.Function): + """ + Mixed-mode (Mix) Causal Convolution Implementation - Combining Triton Forward and CUDA Backward Propagation + + This class implements forward propagation using FLA's Triton kernel, while using the optimized + implementation from TriDao's causal_conv1d CUDA package for backward propagation. + This hybrid strategy combines the advantages of both technologies: + + - Forward: Uses FLA's Triton implementation, optimized for the FLA framework + - Backward: Uses TriDao's causal_conv1d_bwd_function CUDA implementation for faster speed + + Performance Benefits: + - CUDA backward implementation is typically faster than the Triton version, reducing training time + - Maintains the flexibility and compatibility of forward propagation + + Note: + - Input/Output format is (batch, seqlen, dim) + - Backward propagation requires causal_conv1d package: pip install causal-conv1d + - Supports SILU/Swish activation functions + - Current limitations (not yet supported): + * output_final_state must be False + * initial_states must be None + * residual must be None + """ + @staticmethod + @input_guard(no_guard_contiguous=["x"]) + def forward( + ctx, + x, + weight, + bias=None, + residual: torch.Tensor | None = None, + initial_states=None, + output_final_state=False, + activation=None, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + seq_idx: torch.LongTensor | None = None, + ): + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + assert output_final_state is False, "output_final_state must be False for FastCausalConv1dFn" + assert initial_states is None, "initial_states must be None for FastCausalConv1dFn" + assert residual is None, "residual must be None for FastCausalConv1dFn" + + bias = bias.contiguous() if bias is not None else None + if cu_seqlens is not None and seq_idx is None: + seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to( + torch.int32).unsqueeze(0) + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + + # Import here to avoid circular dependency + from fla.modules.conv.triton.ops import causal_conv1d_fwd + + ctx.activation = activation in ["silu", "swish"] + out, _ = causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + residual=None, + initial_state=None, + output_final_state=output_final_state, + activation=activation, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + ) + + ctx.save_for_backward(x, weight, bias, seq_idx, initial_states) + ctx.return_final_states = output_final_state + ctx.return_dinitial_states = ( + initial_states is not None and initial_states.requires_grad + ) + return out, None + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, weight, bias, seq_idx, initial_states = ctx.saved_tensors + dx = torch.empty_like(x, memory_format=torch.contiguous_format) + x = rearrange(x, 'b t d -> b d t') + dx = rearrange(dx, 'b t d -> b d t') + dout = rearrange(dout, 'b t d -> b d t') + dfinal_states = args[0] if ctx.return_final_states else None + + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + # The kernel supports passing in a pre-allocated dx (e.g., in case we want to fuse the + # backward of conv1d with the backward of chunk). + # Here we just pass in None and dx will be allocated in the C++ code. + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + dx, + ctx.return_dinitial_states, + ctx.activation, + ) + dx = rearrange(dx, 'b d t -> b t d') + return ( + dx, + dweight, + dbias if bias is not None else None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def fast_causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, + activation: str | None = None, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + seq_idx: torch.LongTensor | None = None, +): + """ + x: (batch, seqlen, dim) + weight: (dim, width) + bias: (dim,) + seq_idx: (batch, seqlen) + initial_states: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1), to be written to + activation: either None or "silu" or "swish" + + out: (batch, seqlen, dim) + """ + assert causal_conv1d_bwd_function is not None, "causal_conv1d_bwd_function is not available" + return FastCausalConv1dFn.apply( + x, + weight, + bias, + residual, + initial_state, + output_final_state, + activation, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + seq_idx, + ) + + +def causal_conv1d_cuda( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, + activation: str | None = None, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + **kwargs, +): + assert causal_conv1d_fn_cuda is not None, "causal_conv1d_fn_cuda is not available" + seq_idx = kwargs.get('seq_idx') + if cu_seqlens is not None or seq_idx is not None: + assert initial_state is None, "For CUDA backend, initial_state must be None if cu_seqlens or seq_idx is provided" + W = weight.shape[-1] + if x.stride(-1) != 1: + x = x.contiguous() + x_conv1d = rearrange(x, 'b t d -> b d t') + if cu_seqlens is not None and seq_idx is None: + seq_idx = prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu).to(torch.int32).unsqueeze(0) + + y = causal_conv1d_fn_cuda( + x=x_conv1d, + weight=weight, + bias=bias, + activation=activation, + seq_idx=seq_idx, + initial_states=None, + return_final_states=False, + ) + + y = rearrange(y, 'b d t -> b t d') + if output_final_state: + final_state = causal_conv1d_update_states( + x=x, + state_len=W, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ) + else: + final_state = None + if residual is not None: + y.add_(residual) + + return y, final_state diff --git a/fla/modules/conv/long_conv.py b/fla/modules/conv/long_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..ebca73930f372ef342ffa11471a2aea0ec759e70 --- /dev/null +++ b/fla/modules/conv/long_conv.py @@ -0,0 +1,165 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None): + seqlen = u.shape[-1] + fft_size = 2 * seqlen + k_f = torch.fft.rfft(k, n=fft_size) / fft_size + if k_rev is not None: + k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size + k_f = k_f + k_rev_f.conj() + u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size) + + if len(u.shape) > 3: + k_f = k_f.unsqueeze(1) + y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen] + + out = y + u + if gelu: + out = F.gelu(out) + if dropout_mask is not None: + return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype) + else: + return out.to(dtype=u.dtype) + + +class LongConvolution(nn.Module): + """ + LongConvolution applies a convolution operation on the input tensor using a fixed + filter of length max_len. + The filter is learned during training and is applied using FFT convolution. + + Args: + hidden_size (int): The number of expected features in the input and output. + max_len (int): The maximum sequence length. + + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + + def __init__( + self, + hidden_size: int, + max_len: int, + **kwargs, + ): + """ + Initializes the LongConvolution module. + Args: + hidden_size (int): The number of expected features in the input and output. + max_len (int): The maximum sequence length. + """ + super().__init__() + self.hidden_size = hidden_size + self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True) + + def forward(self, x: torch.Tensor, *args, **kwargs): + """ + Applies the LongConvolution operation on the input tensor. + Args: + x: [batch_size, seq_len, hidden_size] tensor + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + x = x.transpose(1, 2) + y = fft_conv(x, self.filter, dropout_mask=None, gelu=False) + y = y.transpose(1, 2) + return y.to(dtype=x.dtype) + + +class PositionalEmbedding(nn.Module): + def __init__(self, emb_dim: int, seq_len: int, **kwargs): + """Complex exponential positional embeddings for implicit long convolution filters.""" + super().__init__() + + self.seq_len = seq_len + # The time embedding fed to the filteres is normalized so that t_f = 1 + t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1 + + if emb_dim > 1: + bands = (emb_dim - 1) // 2 + # To compute the right embeddings we use the "proper" linspace + t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None] + w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1 + + f = torch.linspace(1e-4, bands - 1, bands)[None, None] + z = torch.exp(-1j * f * w) + z = torch.cat([t, z.real, z.imag], dim=-1) + self.z = nn.Parameter(z, requires_grad=False) + + def forward(self, L): + return self.z[:, :L] + + +class ImplicitLongConvolution(nn.Module): + """ + Long convolution with implicit filter parameterized by an MLP. + + Args: + hidden_size (int): + The number of expected features in the input and output. + max_len (int): + The maximum sequence length. + d_emb (Optional[int]): + The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine). + Defaults to 3. + d_hidden (Optional[int]): + The number of features in the hidden layer of the MLP. Defaults to 16. + + Attributes: + pos_emb (`PositionalEmbedding`): The positional embedding layer. + mlp (`nn.Sequential`): The MLP that parameterizes the implicit filter. + + """ + + def __init__( + self, + hidden_size: int, + max_len: int, + d_emb: int = 3, + d_hidden: int = 16, + **kwargs, + ): + """ + Long convolution with implicit filter parameterized by an MLP. + + + """ + super().__init__() + self.hidden_size = hidden_size + self.d_emb = d_emb + + assert ( + d_emb % 2 != 0 and d_emb >= 3 + ), "d_emb must be odd and greater or equal to 3 (time, sine and cosine)" + self.pos_emb = PositionalEmbedding(d_emb, max_len) + + # final linear layer + self.mlp = nn.Sequential( + nn.Linear(d_emb, d_hidden), + torch.nn.ReLU(), + nn.Linear(d_hidden, hidden_size), + ) + + def filter(self, seq_len: int, *args, **kwargs): + return self.mlp(self.pos_emb(seq_len)).transpose(1, 2) + + def forward(self, x: torch.Tensor, *args, **kwargs): + """ + Args: + x: [batch_size, seq_len, hidden_size] tensor + + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + x = x.transpose(1, 2) + k = self.filter(x.shape[-1]) + y = fft_conv(x, k, dropout_mask=None, gelu=False) + + y = y.transpose(1, 2) + return y.to(dtype=x.dtype) diff --git a/fla/modules/conv/short_conv.py b/fla/modules/conv/short_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..ff294172ed9cbd4bc2f07867c94c53d242a5c8b0 --- /dev/null +++ b/fla/modules/conv/short_conv.py @@ -0,0 +1,241 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +"""Short convolution implementation for efficient causal convolutions.""" + +import warnings + +import torch +import torch.nn as nn +from einops import rearrange + +try: + from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda + from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda +except ImportError: + causal_conv1d_fn_cuda = None + causal_conv1d_update_cuda = None + + +class ShortConvolution(nn.Conv1d): + """Short convolution layer for efficient causal convolution operations. + + This class implements a depthwise 1D convolution with causal padding, + designed for efficient sequence processing. It supports multiple backends (Triton/CUDA) + and optional activation functions. + + Args: + hidden_size (int): Number of input/output channels (must be equal for depthwise conv) + kernel_size (int): Size of the convolution kernel + bias (bool, optional): Whether to include learnable bias. Defaults to False. + activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'. + backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'. + device (Optional[torch.device], optional): Device to place the layer on. Defaults to None. + dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None. + **kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility) + + Attributes: + hidden_size (int): Number of channels + activation (Optional[str]): Selected activation function + backend (str): Actual backend being used (may differ from input due to availability) + + Note: + - Uses depthwise convolution (groups=hidden_size) for efficiency + - Applies causal padding (kernel_size-1) to ensure no future information leakage + - Falls back to Triton backend if CUDA backend is unavailable + """ + + def __init__( + self, + hidden_size: int, + kernel_size: int, + bias: bool = False, + activation: str | None = 'silu', + backend: str | None = 'triton', + device: torch.device | None = None, + dtype: torch.dtype | None = None, + **kwargs, + ): + super().__init__( + in_channels=hidden_size, + out_channels=hidden_size, + kernel_size=kernel_size, + groups=hidden_size, + bias=bias, + padding=kernel_size - 1, + device=device, + dtype=dtype, + ) + + self.hidden_size = hidden_size + self.activation = None + + if activation is not None: + assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet." + self.activation = activation + + if 'use_fast_conv1d' in kwargs: + warnings.warn( + "The `use_fast_conv1d` parameter is deprecated and will be ignored. " + "Please use the `backend` parameter instead.", + ) + import os + self.backend = os.environ.get('FLA_CONV_BACKEND', backend) + if backend not in ['cuda', 'triton']: + raise ValueError(f"Invalid backend: {backend}, must be one of ['cuda', 'triton']") + if backend == 'cuda': + if causal_conv1d_fn_cuda is None: + warnings.warn( + "The `backend` parameter is set to `cuda`, but `causal_conv1d_fn` is not available. " + "Switching to the Triton implementation instead. " + "Consider installing `causal_conv1d` to enable the CUDA backend.", + ) + self.backend = 'triton' + + def extra_repr(self): + s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}' + ', stride={stride}') + if self.padding != (0,) * len(self.padding): + s += ', padding={padding}' + if self.dilation != (1,) * len(self.dilation): + s += ', dilation={dilation}' + if self.output_padding != (0,) * len(self.output_padding): + s += ', output_padding={output_padding}' + if self.groups != 1: + s += ', groups={groups}' + if self.bias is None: + s += ', bias=False' + if self.padding_mode != 'zeros': + s += ', padding_mode={padding_mode}' + if self.activation is not None: + s += ', activation={activation}' + s += f', backend={self.backend}' + return s.format(**self.__dict__) + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + cache: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (`torch.Tensor`): + Tensor of shape `[B, T, D]`. `B` must be 1 if `cu_seqlens` is provided. + residual (`Optional[torch.Tensor]`): + Residual tensor of shape `[B, T, D]`. Default: `None`. + mask (`Optional[torch.Tensor]`): + Attention mask dealing with padded positions. + cache (`Optional[torch.Tensor]`): + Previous cache tensor of shape `[N, D, W]`, where `W` is the kernel size. + If provided, the cache is updated **inplace**. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, D, W]`. Default: `False`. + cu_seqlens (Optional[torch.LongTensor]): + Cumulative sequence lengths for each batch. Used for varlen. Default: `None`. + Shape: [B+1] + chunk_indices (Optional[torch.LongTensor]): + Chunk indices for variable-length sequences. Default: `None`. + + Returns: + Tensor of shape `[B, T, D]`. + """ + # Import here to avoid circular dependency + from fla.modules.conv.causal_conv1d import causal_conv1d + + B, T, *_ = x.shape + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + if mask is not None: + if cu_seqlens is not None: + raise ValueError("`mask` and `cu_seqlens` cannot be provided at the same time") + x = x.mul_(mask.unsqueeze(-1)) + + # in decoding phase, the cache (if provided) is updated inplace + if B * T == N: + y, cache = self.step( + x=x, + residual=residual, + cache=cache, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return y, cache + + # cuda backend do not support: + # 1. both `cu_seqlens` and `cache` being provided + # 2. both `cu_seqlens` and `output_final_state` being provided + # and other small issues + # to simplify the implementation, we just switch to triton backend + if self.backend == 'cuda' and cache is not None: + warnings.warn( + "The CUDA backend does not support both `cu_seqlens` and `cache` being provided, " + "or both `cu_seqlens` and `output_final_state` being provided. " + "Switching to the Triton backend instead. ", + stacklevel=2, + ) + self.backend = 'triton' + + return causal_conv1d( + x=x, + weight=rearrange(self.weight, "d 1 w -> d w"), + bias=self.bias, + residual=residual, + initial_state=cache, + output_final_state=output_final_state, + activation=self.activation, + backend=self.backend, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + **kwargs, + ) + + def step( + self, + x: torch.Tensor, + residual: torch.Tensor, + cache: torch.Tensor, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + from fla.modules.conv.triton.ops import causal_conv1d_update + + B, _, D, W = *x.shape, self.kernel_size[0] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + if output_final_state and cache is None: + cache = x.new_zeros(N, D, W) + # NOTE: we follow the fast mode that updates the cache in-place + if self.backend == 'triton': + return causal_conv1d_update( + x=x, + cache=cache, + residual=residual, + weight=rearrange(self.weight, "d 1 w -> d w"), + bias=self.bias, + activation=self.activation, + ) + + shape = x.shape + x = x.squeeze(0) if cu_seqlens is not None else x.squeeze(1) + # equivalent to: + # cache.copy_(cache.roll(shifts=-1, dims=-1)) + # cache[:, :, -1] = x + # y = torch.sum(cache * rearrange(self.weight, "d 1 w -> d w"), dim=-1) + y = causal_conv1d_update_cuda( + x=x, + conv_state=cache, + weight=rearrange(self.weight, "d 1 w -> d w"), + bias=self.bias, + activation=self.activation, + ) + y = y.view(shape) + if residual is not None: + y.add_(residual) + return y, cache + + @property + def state_size(self) -> int: + return self.hidden_size * self.kernel_size diff --git a/fla/modules/conv/triton/__init__.py b/fla/modules/conv/triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cf22a77b0d8b5b8ebd1913895bf0340415f16425 --- /dev/null +++ b/fla/modules/conv/triton/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from .ops import ( + CausalConv1dFunction, + causal_conv1d_bwd, + causal_conv1d_fwd, + causal_conv1d_update, + causal_conv1d_update_states, + compute_dh0_triton, +) + +__all__ = [ + 'CausalConv1dFunction', + 'causal_conv1d_bwd', + 'causal_conv1d_fwd', + 'causal_conv1d_update', + 'causal_conv1d_update_states', + 'compute_dh0_triton', +] diff --git a/fla/modules/conv/triton/kernels.py b/fla/modules/conv/triton/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..fd351fabea43728848cf127d7e60940fc88713ee --- /dev/null +++ b/fla/modules/conv/triton/kernels.py @@ -0,0 +1,657 @@ +import torch +import triton +import triton.language as tl +from einops import rearrange + +from fla.utils import IS_AMD, autotune_cache_kwargs, input_guard + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] +STATIC_WARPS = 32 if not IS_AMD else 16 + + +@triton.heuristics({ + 'HAS_WEIGHT': lambda args: args['weight'] is not None, + 'HAS_BIAS': lambda args: args['bias'] is not None, + 'HAS_RESIDUAL': lambda args: args['residual'] is not None, + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [16, 32, 64, 128] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D', 'W', 'NB'], + **autotune_cache_kwargs, +) +@triton.jit +def causal_conv1d_fwd_kernel( + x, + y, + weight, + bias, + residual, + cu_seqlens, + initial_state, + chunk_indices, + B, + T, + stride_x_n, + stride_x_t, + stride_x_d, + D: tl.constexpr, + W: tl.constexpr, + BT: tl.constexpr, + BW: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + ACTIVATION: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + p_x = x + bos * stride_x_t + else: + i_n = i_b + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + p_x = x + tl.cast(i_b, tl.int64) * stride_x_n + + o_d = i_d * BD + tl.arange(0, BD) + o_w = tl.arange(0, BW) + W - BW + m_d = o_d < D + m_w = o_w >= 0 + + if HAS_WEIGHT: + # [BD, BW] + b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32) + + b_y = tl.zeros((BT, BD), dtype=tl.float32) + if not USE_INITIAL_STATE: + for i_w in tl.static_range(-W + 1, 1): + p_yi = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + # [BT, BD] + b_yi = tl.load(p_yi, boundary_check=(0, 1)).to(tl.float32) + if HAS_WEIGHT: + b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) + b_y += b_yi + elif i_t * BT >= W: + # to make Triton compiler happy, we need to copy codes + for i_w in tl.static_range(-W + 1, 1): + p_yi = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + # [BT, BD] + b_yi = tl.load(p_yi, boundary_check=(0, 1)).to(tl.float32) + if HAS_WEIGHT: + b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) + b_y += b_yi + else: + o_t = i_t * BT + tl.arange(0, BT) + for i_w in tl.static_range(-W + 1, 1): + o_x = o_t + i_w + m_x = ((o_x >= 0) & (o_x < T))[:, None] & m_d + m_c = ((o_x + W >= 0) & (o_x < 0))[:, None] & m_d + + b_yi = tl.load( + p_x + o_x[:, None] * stride_x_t + o_d * stride_x_d, + mask=m_x, + other=0 + ).to(tl.float32) + + b_yi += tl.load(initial_state + i_n * D*W + o_d * W + (o_x + W)[:, None], mask=m_c, other=0).to(tl.float32) + + if HAS_WEIGHT: + b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) + b_y += b_yi + + if HAS_BIAS: + b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32) + + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + b_y = b_y * tl.sigmoid(b_y) + + if HAS_RESIDUAL: + p_residual = tl.make_block_ptr(residual + bos * D, (T, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + b_residual = tl.load(p_residual, boundary_check=(0, 1)) + b_y += b_residual + + p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + tl.store(p_y, tl.cast(b_y, dtype=p_y.dtype.element_ty, fp_downcast_rounding='rtne'), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'HAS_WEIGHT': lambda args: args['dw'] is not None, + 'HAS_BIAS': lambda args: args['db'] is not None, + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'USE_FINAL_STATE': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [16, 32, 64, 128] + for num_warps in [4, 8, 16, 32] + ], + key=['D', 'W', 'NB'], + **autotune_cache_kwargs, +) +@triton.jit +def causal_conv1d_bwd_kernel( + x, + y, + weight, + initial_state, + dht, + dy, + dx, + dw, + db, + cu_seqlens, + chunk_indices, + B, + T, + stride_x_n, # x batch stride + stride_x_t, # x time stride + stride_x_d, # x dim stride + stride_dx_n, # dx batch stride + stride_dx_t, # dx time stride + stride_dx_d, # dx dim stride + D: tl.constexpr, + W: tl.constexpr, + BT: tl.constexpr, + BW: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + ACTIVATION: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + p_x = x + bos * stride_x_t + else: + i_tg = i_b * tl.num_programs(1) + i_t + i_n = i_b + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + p_x = x + tl.cast(i_b, tl.int64) * stride_x_n + + o_d = i_d * BD + tl.arange(0, BD) + o_w = tl.arange(0, BW) + W - BW + m_d = o_d < D + m_w = o_w >= 0 + + if HAS_WEIGHT: + p_x = tl.make_block_ptr(p_x, (T, D), (stride_x_t, stride_x_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)) + # [BD, BW] + b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0) + + b_dx = tl.zeros((BT, BD), dtype=tl.float32) + if HAS_BIAS: + b_db = tl.zeros((BD,), dtype=tl.float32) + + if not USE_FINAL_STATE and not USE_INITIAL_STATE: + for i_w in tl.static_range(0, W): + p_dy = tl.make_block_ptr(dy + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + # [BT, BD] + b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) + b_ys = tl.sigmoid(b_y) + b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys)) + b_wdy = b_dy + if HAS_WEIGHT: + # [BT, BD] + b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) + # [BD] + b_dw = tl.sum(b_dy * b_x, 0) + tl.store(dw + i_tg * D*W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d) + if HAS_BIAS and i_w == 0: + b_db += tl.sum(b_dy, 0) + b_dx += b_wdy + elif i_t * BT >= W: + # to make Triton compiler happy, we need to copy codes + for i_w in tl.static_range(0, W): + p_dy = tl.make_block_ptr(dy + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + # [BT, BD] + b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) + b_ys = tl.sigmoid(b_y) + b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys)) + b_wdy = b_dy + if HAS_WEIGHT: + # [BT, BD] + b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) + # [BD] + b_dw = tl.sum(b_dy * b_x, 0) + tl.store(dw + i_tg * D*W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d) + if HAS_BIAS and i_w == 0: + b_db += tl.sum(b_dy, 0) + b_dx += b_wdy + else: + # which may use initial state + o_t = i_t * BT + tl.arange(0, BT) + for i_w in tl.static_range(0, W): + p_dy = tl.make_block_ptr(dy + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + b_dy_shift = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + p_y = tl.make_block_ptr(y + bos * D, (T, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0)) + b_y_shift = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) + b_ys = tl.sigmoid(b_y_shift) + b_dy_shift = b_dy_shift * b_ys * (1 + b_y_shift * (1 - b_ys)) + if HAS_WEIGHT: + # gradient comes from x：sum_t dy[t+i_w] * x[t] + b_dw = tl.sum(b_dy_shift * b_x, 0) + # index of cache：c = W - i_w + t + if USE_INITIAL_STATE: + mask_head_rows = (o_t < i_w) & (o_t < T) + # dy_head = dy[t] + b_dy_head = tl.load(dy + bos * D + o_t[:, None] * D + o_d, mask=(mask_head_rows[:, None] & m_d[None, :]), + other=0.0).to(tl.float32) + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + # use y[t] （not y[t+i_w]） + b_y_head = tl.load(y + bos * D + o_t[:, None] * D + o_d, + mask=(mask_head_rows[:, None] & m_d[None, :]), other=0.0).to(tl.float32) + b_ys_head = tl.sigmoid(b_y_head) + b_dy_head = b_dy_head * b_ys_head * (1 + b_y_head * (1 - b_ys_head)) + o_c = W - i_w + o_t + # index 0 is padding 0 + mask_c = (mask_head_rows & (o_c >= 1) & (o_c < W)) + b_xc = tl.load(initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None], + mask=(mask_c[:, None] & m_d[None, :]), other=0.0).to(tl.float32) + # add the gradient comes from initial_state + b_dw += tl.sum(b_dy_head * b_xc, 0) + tl.store(dw + i_tg * D * W + o_d * W + W - i_w - 1, b_dw.to(dw.dtype.element_ty), mask=m_d) + + if HAS_BIAS and i_w == 0: + b_db += tl.sum(b_dy_shift, 0) + b_wdy = b_dy_shift if not HAS_WEIGHT else (b_dy_shift * tl.sum(b_w * (o_w == (W - i_w - 1)), 1)) + b_dx += b_wdy + + if HAS_BIAS: + b_db = tl.cast(b_db, dtype=db.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(db + i_tg * D + o_d, b_db, mask=m_d) + + if USE_FINAL_STATE: + if i_t * BT + BT >= T-W: + start_tok = max(0, T - (W - 1)) + offset = i_t * BT + tl.arange(0, BT) + tok_idx = offset - start_tok + mask = (offset >= start_tok) & (offset < T) + w_idx = 1 + tok_idx + dht_off = i_n * D * W + o_d[None, :] * W + w_idx[:, None] + b_dht = tl.load(dht + dht_off, mask=mask[:, None] & m_d[None, :], other=0.).to(tl.float32) + b_dx += b_dht + + if IS_VARLEN: + p_dx = dx + bos * stride_dx_t + else: + p_dx = dx + tl.cast(i_b, tl.int64) * stride_dx_n + + p_dx = tl.make_block_ptr(p_dx, (T, D), (stride_dx_t, stride_dx_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + tl.store(p_dx, tl.cast(b_dx, dtype=p_dx.dtype.element_ty, fp_downcast_rounding='rtne'), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, + 'HAS_WEIGHT': lambda args: args['weight'] is not None, + 'HAS_BIAS': lambda args: args['bias'] is not None, + 'HAS_RESIDUAL': lambda args: args['residual'] is not None, +}) +@triton.jit +def causal_conv1d_update_kernel( + x, + cache, + residual, + y, + weight, + bias, + stride_x_n, # batch stride + stride_x_d, # dim stride + stride_y_n, # batch stride + stride_y_d, # dim stride + D: tl.constexpr, + W: tl.constexpr, + BD: tl.constexpr, + BW: tl.constexpr, + ACTIVATION: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, +): + i_d, i_n = tl.program_id(0), tl.program_id(1) + + o_d = i_d * BD + tl.arange(0, BD) + o_w = tl.arange(0, BW) + m_d = o_d < D + m_w = o_w < W + + # [BD] + b_x = tl.load(x + i_n * stride_x_n + o_d * stride_x_d, mask=m_d, other=0).to(tl.float32) + + b_cache = tl.zeros((BD, BW), dtype=tl.float32) + + if USE_INITIAL_STATE: + # 2. Shift Cache (Read [1:]) + p_cache_read = tl.make_block_ptr( + cache + i_n * D*W, + shape=(D, W), + strides=(W, 1), + offsets=(i_d * BD, 1), + block_shape=(BD, BW), + order=(1, 0) + ) + b_cache = tl.load(p_cache_read, boundary_check=(0, 1)).to(tl.float32) + + # 3. Fill x to the last position + m_update = o_w == (W - 1) + b_cache = tl.where(m_update[None, :], b_x[:, None], b_cache) + + if HAS_WEIGHT: + b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0) + b_y = tl.sum(b_cache * b_w, 1) + else: + b_y = tl.sum(b_cache, 1) + + if HAS_BIAS: + b_y += tl.load(bias + o_d, mask=m_d) + + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + b_y = b_y * tl.sigmoid(b_y) + + if HAS_RESIDUAL: + b_y += tl.load(residual + i_n * D + o_d, mask=m_d, other=0) + + tl.store(y + i_n * stride_y_n + o_d * stride_y_d, tl.cast(b_y, + dtype=y.dtype.element_ty, fp_downcast_rounding='rtne'), mask=m_d) + + if USE_INITIAL_STATE: + p_cache_write = tl.make_block_ptr( + cache + i_n * D*W, + shape=(D, W), + strides=(W, 1), + offsets=(i_d * BD, 0), + block_shape=(BD, BW), + order=(1, 0) + ) + tl.store(p_cache_write, tl.cast(b_cache, dtype=cache.dtype.element_ty, + fp_downcast_rounding='rtne'), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_ACTIVATION': lambda args: args['y'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit +def compute_dh0_kernel( + dy, + y, + weight, + dh0, + cu_seqlens, + stride_dy_n, + stride_dy_t, + T, + D: tl.constexpr, + W: tl.constexpr, + BD: tl.constexpr, + USE_ACTIVATION: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Compute dh0 (gradient w.r.t. initial_state) in a separate kernel. + This avoids Triton compiler bugs on some architectures (e.g., GB200). + + Grid: (cdiv(D, BD), N) + """ + i_d, i_n = tl.program_id(0), tl.program_id(1) + + # Get sequence boundaries + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + seq_len = eos - bos + # For varlen, dy is [1, total_T, D], offset by bos + dy_base = dy + bos * stride_dy_t + else: + seq_len = T + # For non-varlen, dy is [B, T, D], offset by i_n * stride_dy_n + dy_base = dy + tl.cast(i_n, tl.int64) * stride_dy_n + + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + + # For each i_w in [1, W), compute dh0[i_n, :, i_w] + for i_w in tl.static_range(1, W): + b_dh0 = tl.zeros([BD], dtype=tl.float32) + + # Accumulate contributions from t = 0 to min(i_w, seq_len) - 1 + for t in tl.static_range(0, W - 1): + if t < i_w: + w_idx = i_w - 1 - t + + # Load dy[t, :] relative to dy_base + p_dy = dy_base + t * stride_dy_t + o_d + m_t = (t < seq_len) & m_d + b_dy = tl.load(p_dy, mask=m_t, other=0).to(tl.float32) + + if USE_ACTIVATION: + if IS_VARLEN: + p_y = y + bos * stride_dy_t + t * stride_dy_t + o_d + else: + p_y = y + tl.cast(i_n, tl.int64) * stride_dy_n + t * stride_dy_t + o_d + b_y = tl.load(p_y, mask=m_t, other=0).to(tl.float32) + b_ys = tl.sigmoid(b_y) + b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys)) + + # Get weight[:, w_idx] + b_w_col = tl.load(weight + o_d * W + w_idx, mask=m_d, other=0).to(tl.float32) + + # Accumulate + b_dh0 += tl.where(m_t, b_dy * b_w_col, 0) + + # Store dh0[i_n, :, i_w] + p_dh0 = dh0 + i_n * D * W + o_d * W + i_w + tl.store(p_dh0, b_dh0.to(dh0.dtype.element_ty), mask=m_d) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit +def causal_conv1d_states_fwd_kernel( + x, + initial_state, + final_state, + cu_seqlens, + T, + D, + W, + stride_x_n, + stride_x_t, + stride_x_d, + BD: tl.constexpr, + BW: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_n = tl.program_id(0), tl.program_id(1) + + # o_d Shape: [BD] + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + seq_len = (eos - bos).to(tl.int32) + p_x = x + bos * stride_x_t + else: + seq_len = T + p_x = x + tl.cast(i_n, tl.int64) * stride_x_n + + p_x = tl.make_block_ptr(p_x, (seq_len, D), (stride_x_t, stride_x_d), (seq_len - BW, i_d * BD), (BW, BD), (1, 0)) + + # b_x Shape: [BW, BD] + b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + + if USE_INITIAL_STATE: + if seq_len < BW: + o_c = W - (BW - seq_len) + tl.arange(0, BW) + m_c = (o_c >= 0) & (o_c < W) + + p_init = initial_state + i_n * D*W + o_d[None, :] * W + o_c[:, None] + mask_init = m_d[None, :] & m_c[:, None] + + b_cache = tl.load(p_init, mask=mask_init, other=0) + b_x += b_cache + + # final_state: [N, D, W] (Channel Major inside sample) + # o_w Shape: [BW] + o_w = W - BW + tl.arange(0, BW) + + # o_d[:, None] -> [BD, 1] + # o_w[None, :] -> [1, BW] + # p_final Shape -> [BD, BW] + p_final = final_state + tl.cast(i_n, tl.int64) * D*W + o_d[:, None] * W + o_w[None, :] + + # m_final Shape -> [BD, BW] + m_final = m_d[:, None] & (o_w[None, :] >= 0) + + tl.store(p_final, tl.trans(b_x).to(final_state.dtype.element_ty), mask=m_final) + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d_update_states( + x: torch.Tensor, + state_len: int, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> torch.Tensor: + if cu_seqlens is not None: + N = len(cu_seqlens) - 1 + if x.dim() == 2: + stride_x_n = 0 + stride_x_t, stride_x_d = x.stride() + T = x.shape[0] + else: + stride_x_n = x.stride(0) + stride_x_t, stride_x_d = x.stride(1), x.stride(2) + T = x.shape[1] + D = x.shape[-1] + else: + B, T, D = x.shape + N = B + stride_x_n, stride_x_t, stride_x_d = x.stride() + + W = state_len + final_state = torch.empty(N, D, W, dtype=x.dtype, device=x.device) + + BD = min(triton.next_power_of_2(D), 256) + BW = triton.next_power_of_2(W) + + grid = (triton.cdiv(D, BD), N) + + causal_conv1d_states_fwd_kernel[grid]( + x=x, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + T=T, + D=D, + W=W, + stride_x_n=stride_x_n, + stride_x_t=stride_x_t, + stride_x_d=stride_x_d, + BW=BW, + BD=BD, + ) + return final_state + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d_update( + x: torch.Tensor, + cache: torch.Tensor, + residual: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + activation: str | None = None, +) -> torch.Tensor: + shape = x.shape + if weight is not None and x.shape[-1] != weight.shape[0]: + x = rearrange(x, 'b t ... -> b t (...)') + + D = x.shape[-1] + N = x.numel() // D + W = weight.shape[1] if weight is not None else None + BD = 8 + BW = triton.next_power_of_2(W) + + if x.dim() == 2: + # Case: (N, D) + stride_x_n = x.stride(0) + stride_x_d = x.stride(1) + elif x.dim() == 3 and x.shape[0] == 1: + # Case: (1, N, D) -> Time=1, Batch=N, Dim=D + # Batch 在 dim 1 + stride_x_n = x.stride(1) + stride_x_d = x.stride(2) + elif x.dim() == 3: + # Case: (N, 1, D) -> Batch=N, Time=1, Dim=D + # Batch 在 dim 0 + stride_x_n = x.stride(0) + stride_x_d = x.stride(2) + else: + # Fallback / Error case + raise ValueError(f"Unsupported input shape: {x.shape}") + + y = torch.empty_like(x, memory_format=torch.contiguous_format) + + if y.dim() == 2: + stride_y_n, stride_y_d = y.stride(0), y.stride(1) + elif y.dim() == 3 and y.shape[0] == 1: + stride_y_n, stride_y_d = y.stride(1), y.stride(2) + elif y.dim() == 3: + stride_y_n, stride_y_d = y.stride(0), y.stride(2) + + def grid(meta): return (triton.cdiv(D, meta['BD']), N) + + causal_conv1d_update_kernel[grid]( + x=x, + cache=cache, + residual=residual, + y=y, + weight=weight, + bias=bias, + stride_x_n=stride_x_n, + stride_x_d=stride_x_d, + stride_y_n=stride_y_n, + stride_y_d=stride_y_d, + D=D, + W=W, + BD=BD, + BW=BW, + ACTIVATION=activation, + num_warps=STATIC_WARPS, + ) + return y.view(shape), cache diff --git a/fla/modules/conv/triton/ops.py b/fla/modules/conv/triton/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed6dd64684732ddfad0f9c4a9a3d0351f05ae87 --- /dev/null +++ b/fla/modules/conv/triton/ops.py @@ -0,0 +1,397 @@ + +import torch +import triton +from einops import rearrange + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import input_guard + +from .kernels import ( + STATIC_WARPS, + causal_conv1d_bwd_kernel, + causal_conv1d_fwd_kernel, + causal_conv1d_states_fwd_kernel, + causal_conv1d_update_kernel, + compute_dh0_kernel, +) + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + activation: str | None = None, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + BT: int = 64, +) -> torch.Tensor: + shape = x.shape + if x.shape[-1] != weight.shape[0]: + x = rearrange(x, 'b t ... -> b t (...)') + B, T, D = x.shape[0], x.shape[1], weight.shape[0] + W = weight.shape[1] + stride_x_n, stride_x_t, stride_x_d = x.stride() + + BW = triton.next_power_of_2(W) + if cu_seqlens is not None and chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + NB = triton.cdiv(B*T, 1024) + + y = torch.empty_like(x, memory_format=torch.contiguous_format) + + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B) + causal_conv1d_fwd_kernel[grid]( + x=x, + y=y, + weight=weight, + bias=bias, + residual=residual, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + chunk_indices=chunk_indices, + B=B, + T=T, + D=D, + W=W, + BT=BT, + BW=BW, + NB=NB, + stride_x_n=stride_x_n, + stride_x_t=stride_x_t, + stride_x_d=stride_x_d, + ACTIVATION=activation, + ) + final_state = None + if output_final_state: + final_state = causal_conv1d_update_states( + x=x, + state_len=W, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ) + return y.view(shape), final_state + + +def compute_dh0_triton( + dy: torch.Tensor, + y: torch.Tensor | None, + weight: torch.Tensor, + initial_state: torch.Tensor, + activation: str | None, + cu_seqlens: torch.Tensor | None, +) -> torch.Tensor: + """ + Compute dh0 (gradient w.r.t. initial_state) using a separate Triton kernel. + This is a workaround for Triton compiler bugs on some architectures (e.g., GB200). + """ + D, W = weight.shape + N = initial_state.shape[0] + T = dy.shape[1] + + # Initialize dh0 + dh0 = torch.zeros_like(initial_state) + + BD = 32 + grid = (triton.cdiv(D, BD), N) + + y_to_pass = y if activation in ('swish', 'silu') else None + # dy is [B, T, D], stride_n = T*D, stride_t = D + stride_dy_n = dy.stride(0) + stride_dy_t = dy.stride(1) + + compute_dh0_kernel[grid]( + dy=dy, + y=y_to_pass, + weight=weight, + dh0=dh0, + cu_seqlens=cu_seqlens, + stride_dy_n=stride_dy_n, + stride_dy_t=stride_dy_t, + T=T, + D=D, + W=W, + BD=BD, + ) + + return dh0 + + +def causal_conv1d_bwd( + x: torch.Tensor, + dy: torch.Tensor, + dht: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + activation: str | None = None, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + BT: int = 64, +): + shape = x.shape + if x.shape[-1] != weight.shape[0]: + x = rearrange(x, 'b t ... -> b t (...)') + B, T, D = x.shape + W = weight.shape[1] if weight is not None else None + + stride_x_n, stride_x_t, stride_x_d = x.stride() + + BW = triton.next_power_of_2(W) + if cu_seqlens is not None and chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + NB = triton.cdiv(B*T, 1024) + + y = None + if activation is not None: + y, _ = causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + residual=None, + initial_state=initial_state, + activation=None, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + output_final_state=False, + chunk_indices=chunk_indices, + ) + dx = torch.empty_like(x) + dw = weight.new_empty(B*NT, *weight.shape, dtype=torch.float) if weight is not None else None + db = bias.new_empty(B*NT, *bias.shape, dtype=torch.float) if bias is not None else None + dr = dy if residual is not None else None + + stride_dx_n, stride_dx_t, stride_dx_d = dx.stride() + + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B) + causal_conv1d_bwd_kernel[grid]( + x=x, + y=y, + weight=weight, + initial_state=initial_state, + dht=dht, + dy=dy, + dx=dx, + dw=dw, + db=db, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + D=D, + W=W, + BT=BT, + BW=BW, + NB=NB, + stride_x_n=stride_x_n, + stride_x_t=stride_x_t, + stride_x_d=stride_x_d, + stride_dx_n=stride_dx_n, + stride_dx_t=stride_dx_t, + stride_dx_d=stride_dx_d, + ACTIVATION=activation, + ) + if weight is not None: + dw = dw.sum(0).to(weight) + if bias is not None: + db = db.sum(0).to(bias) + + # Compute dh0 using separate Triton kernel to avoid compiler bugs on some architectures (e.g., GB200) + dh0 = None + if initial_state is not None: + dh0 = compute_dh0_triton( + dy=dy, + y=y, + weight=weight, + initial_state=initial_state, + activation=activation, + cu_seqlens=cu_seqlens, + ) + + return dx.view(shape), dw, db, dr, dh0 + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d_update_states( + x: torch.Tensor, + state_len: int, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> torch.Tensor: + if cu_seqlens is not None: + N = len(cu_seqlens) - 1 + if x.dim() == 2: + stride_x_n = 0 + stride_x_t, stride_x_d = x.stride() + T = x.shape[0] + else: + stride_x_n = x.stride(0) + stride_x_t, stride_x_d = x.stride(1), x.stride(2) + T = x.shape[1] + D = x.shape[-1] + else: + B, T, D = x.shape + N = B + stride_x_n, stride_x_t, stride_x_d = x.stride() + + W = state_len + final_state = torch.empty(N, D, W, dtype=x.dtype, device=x.device) + + BD = min(triton.next_power_of_2(D), 256) + BW = triton.next_power_of_2(W) + + grid = (triton.cdiv(D, BD), N) + + causal_conv1d_states_fwd_kernel[grid]( + x=x, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + T=T, + D=D, + W=W, + stride_x_n=stride_x_n, + stride_x_t=stride_x_t, + stride_x_d=stride_x_d, + BW=BW, + BD=BD, + ) + return final_state + + +@input_guard(no_guard_contiguous=["x"]) +def causal_conv1d_update( + x: torch.Tensor, + cache: torch.Tensor, + residual: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + activation: str | None = None, +) -> torch.Tensor: + shape = x.shape + if weight is not None and x.shape[-1] != weight.shape[0]: + x = rearrange(x, 'b t ... -> b t (...)') + + D = x.shape[-1] + N = x.numel() // D + W = weight.shape[1] if weight is not None else None + BD = 8 + BW = triton.next_power_of_2(W) + + if x.dim() == 2: + # Case: (N, D) + stride_x_n = x.stride(0) + stride_x_d = x.stride(1) + elif x.dim() == 3 and x.shape[0] == 1: + # Case: (1, N, D) -> Time=1, Batch=N, Dim=D + # Batch 在 dim 1 + stride_x_n = x.stride(1) + stride_x_d = x.stride(2) + elif x.dim() == 3: + # Case: (N, 1, D) -> Batch=N, Time=1, Dim=D + # Batch 在 dim 0 + stride_x_n = x.stride(0) + stride_x_d = x.stride(2) + else: + # Fallback / Error case + raise ValueError(f"Unsupported input shape: {x.shape}") + + y = torch.empty_like(x, memory_format=torch.contiguous_format) + + if y.dim() == 2: + stride_y_n, stride_y_d = y.stride(0), y.stride(1) + elif y.dim() == 3 and y.shape[0] == 1: + stride_y_n, stride_y_d = y.stride(1), y.stride(2) + elif y.dim() == 3: + stride_y_n, stride_y_d = y.stride(0), y.stride(2) + + def grid(meta): return (triton.cdiv(D, meta['BD']), N) + + causal_conv1d_update_kernel[grid]( + x=x, + cache=cache, + residual=residual, + y=y, + weight=weight, + bias=bias, + stride_x_n=stride_x_n, + stride_x_d=stride_x_d, + stride_y_n=stride_y_n, + stride_y_d=stride_y_d, + D=D, + W=W, + BD=BD, + BW=BW, + ACTIVATION=activation, + num_warps=STATIC_WARPS, + ) + return y.view(shape), cache + + +class CausalConv1dFunction(torch.autograd.Function): + + @staticmethod + @input_guard(no_guard_contiguous=["x"]) + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, + activation: str | None = None, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + ): + BT = chunk_size + if cu_seqlens is not None and chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) + ctx.activation = activation + ctx.cu_seqlens = cu_seqlens + ctx.cu_seqlens_cpu = cu_seqlens_cpu + ctx.chunk_indices = chunk_indices + ctx.save_for_backward(x, weight, bias, residual, initial_state) + y, final_state = causal_conv1d_fwd( + x=x, + weight=weight, + bias=bias, + residual=residual, + initial_state=initial_state, + output_final_state=output_final_state, + activation=activation, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + BT=BT, + ) + return y, final_state + + @staticmethod + @input_guard(no_guard_contiguous=["dy"]) + def backward(ctx, dy: torch.Tensor, dht: torch.Tensor | None = None): + x, weight, bias, residual, initial_state = ctx.saved_tensors + dx, dw, db, dr, dh0 = causal_conv1d_bwd( + x=x, + dy=dy, + dht=dht, + weight=weight, + bias=bias, + residual=residual, + initial_state=initial_state, + activation=ctx.activation, + cu_seqlens=ctx.cu_seqlens, + cu_seqlens_cpu=ctx.cu_seqlens_cpu, + chunk_indices=ctx.chunk_indices, + ) + return dx, dw, db, dr, dh0, None, None, None, None, None, None diff --git a/fla/modules/convolution.py b/fla/modules/convolution.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1e9a2ac6acc700feb143824a2c799d91d8ae9c --- /dev/null +++ b/fla/modules/convolution.py @@ -0,0 +1,38 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +from fla.modules.conv import ( + ImplicitLongConvolution, + LongConvolution, + PositionalEmbedding, + ShortConvolution, + causal_conv1d, + fft_conv, +) +from fla.modules.conv.cp import CausalConv1dFunctionCP, causal_conv1d_cp +from fla.modules.conv.cuda import FastCausalConv1dFn, fast_causal_conv1d_fn +from fla.modules.conv.triton import ( + CausalConv1dFunction, + causal_conv1d_bwd, + causal_conv1d_fwd, + causal_conv1d_update, + causal_conv1d_update_states, +) + +__all__ = [ + 'CausalConv1dFunction', + 'CausalConv1dFunctionCP', + 'FastCausalConv1dFn', + 'ImplicitLongConvolution', + 'LongConvolution', + 'PositionalEmbedding', + 'ShortConvolution', + 'causal_conv1d', + 'causal_conv1d_bwd', + 'causal_conv1d_cp', + 'causal_conv1d_fwd', + 'causal_conv1d_update', + 'causal_conv1d_update_states', + 'fast_causal_conv1d_fn', + 'fft_conv', +] diff --git a/fla/modules/feature_map.py b/fla/modules/feature_map.py new file mode 100644 index 0000000000000000000000000000000000000000..15f3b194f0997f7fd5735768b48d3561c2fe6360 --- /dev/null +++ b/fla/modules/feature_map.py @@ -0,0 +1,298 @@ + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from fla.modules.activations import fast_gelu_impl, sigmoid, sqrelu, swish +from fla.modules.layernorm import layer_norm +from fla.utils import checkpoint + + +@checkpoint +def flatten_diag_outer_product(x, y): + z = torch.einsum("...i,...j->...ij", x, y) + N = z.size(-1) + indicies = torch.triu_indices(N, N) + return z[..., indicies[0], indicies[1]] + + +@checkpoint +def flatten_diag_outer_product_off1(x, y): + z = torch.einsum("...i,...j->...ij", x, y) + N = z.size(-1) + indicies = torch.triu_indices(N, N, 1) + indices2 = torch.arange(0, N) + return z[..., indicies[0], indicies[1]], z[..., indices2, indices2] + + +def is_power_of_2(n): + return (n & (n - 1) == 0) and n != 0 + + +class HedgehogFeatureMap(nn.Module): + + r""" + Hedgehog feature map as introduced in + `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry `_ + """ + + def __init__( + self, + head_dim: int, + ) -> HedgehogFeatureMap: + super().__init__() + # Trainable map + self.layer = nn.Linear(head_dim, head_dim) + self.init_weights_() + + def init_weights_(self): + """Initialize trainable map as identity""" + with torch.no_grad(): + identity = torch.eye(*self.layer.weight.shape[-2:], dtype=torch.float) + self.layer.weight.copy_(identity.to(self.layer.weight)) + nn.init.zeros_(self.layer.bias) + + def forward(self, x: torch.Tensor): + x = self.layer(x) # shape b, h, l, d + return torch.cat([2*x, -2*x], dim=-1).softmax(-1) + + +class T2RFeatureMap(nn.Module): + + r""" + Simple linear mapping feature map as in + `Finetuning Pretrained Transformers into RNNs `_ + """ + + def __init__( + self, + head_dim: int, + dot_dim: int = None, + bias: bool | None = False, + ) -> T2RFeatureMap: + super().__init__() + # Trainable map + if dot_dim is None: + dot_dim = head_dim + + self.head_dim = head_dim + self.dot_dim = dot_dim + self.bias = bias + + self.layer = nn.Linear(head_dim, dot_dim, bias=bias) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(head_dim={self.head_dim}, dot_dim={self.dot_dim}, bias={self.bias})" + + def forward(self, x: torch.Tensor): + return self.layer(x).relu() + + +class DPFPFeatureMap(nn.Module): + + r""" + Deterministic Parameter-Free Projection (DPFP) feature map in + `Linear Transformers Are Secretly Fast Weight Programmers `_ + """ + + def __init__( + self, + head_dim: int, + nu: int = 4, + ) -> DPFPFeatureMap: + super().__init__() + self.nu = nu + + def forward(self, x: torch.Tensor): + x = torch.cat([x.relu(), -x.relu()], dim=-1) + x_rolled = torch.cat([x.roll(shifts=j, dims=-1) for j in range(1, self.nu+1)], dim=-1) + x_repeat = torch.cat([x] * self.nu, dim=-1) + return x_repeat * x_rolled + + +class HadamardFeatureMap(nn.Module): + def __init__( + self, + head_dim: int, + ) -> HadamardFeatureMap: + super().__init__() + # Trainable map + self.layer1 = nn.Linear(head_dim, head_dim) + self.layer2 = nn.Linear(head_dim, head_dim) + + def forward(self, x: torch.Tensor): + return self.layer1(x) * self.layer2(x) + + +class LearnableOuterProductFeatureMap(nn.Module): + def __init__( + self, + head_dim: int, + feature_dim: int, + ) -> LearnableOuterProductFeatureMap: + super().__init__() + # Trainable map + self.layer1 = nn.Linear(head_dim, feature_dim, bias=False) + self.layer2 = nn.Linear(head_dim, feature_dim, bias=False) + self.normalizer = feature_dim ** -0.5 + + def forward(self, x: torch.Tensor): + return flatten_diag_outer_product(self.layer1(x), self.layer2(x)) + + +class LearnablePolySketchNonNegativeFeatureMap(nn.Module): + + def __init__( + self, + head_dim: int, + sketch_size: int | None = None, + degree: int | None = 2, + ) -> LearnablePolySketchNonNegativeFeatureMap: + super().__init__() + + assert is_power_of_2(degree) and degree >= 2, f"The degree {degree} must be a power of 2" + + self.head_dim = head_dim + self.sketch_size = sketch_size if sketch_size is not None else head_dim + self.degree = degree + + self.gamma = nn.Parameter(torch.ones(head_dim)) + self.beta = nn.Parameter(torch.zeros(head_dim)) + # NOTE: the sketch layers defined here are quite different from the original paper + # currently we simply use linear layers without any non-linear activations + self.sketches1 = nn.ModuleList([ + nn.Linear(head_dim, sketch_size, bias=False), + *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)], + ]) + self.sketches2 = nn.ModuleList([ + nn.Linear(head_dim, sketch_size, bias=False), + *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)], + ]) + + def forward(self, x: torch.Tensor): + # Section 2.1 + x = layer_norm(x, self.gamma, self.beta) + # first map the input to sketch size with learnable parameters + x = self.sketches1[0](x) * self.sketches2[0](x) * self.head_dim ** -0.5 + for i in range(1, int(math.log2(self.degree)) - 1): + x = self.sketches1[i](x) * self.sketches2[i](x) * self.head_dim ** -0.5 + # do sketch mapping for log2(p) - 1 times in total + # do p=2 mapping to ensure non-negativity + return flatten_diag_outer_product(x, x) + + +class TaylorFeatureMap(nn.Module): + def __init__( + self, + head_dim: int, + ) -> TaylorFeatureMap: + super().__init__() + self.head_dim = head_dim + self.r2 = math.sqrt(2) + self.rd = math.sqrt(self.head_dim) + self.rrd = math.sqrt(self.rd) + + def forward(self, x: torch.Tensor): + x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) + return torch.cat([torch.ones_like(x[..., 0:1]), x / self.rrd, x2_2 / (self.rd * self.r2), x2_1 / self.rd], dim=-1) + + +class RebasedFeatureMap(nn.Module): + + def __init__( + self, + head_dim: int, + use_gamma: bool | None = True, + use_beta: bool | None = True, + normalize: bool | None = True, + ) -> RebasedFeatureMap: + super().__init__() + + self.head_dim = head_dim + self.use_gamma = use_gamma + self.use_beta = use_beta + self.normalize = normalize + + self.gamma = None + self.beta = None + if use_gamma: + self.gamma = nn.Parameter(torch.ones(head_dim)) + if use_beta: + self.beta = nn.Parameter(torch.zeros(head_dim)) + + def forward(self, x: torch.Tensor, flatten: bool | None = True): + if self.use_beta and self.use_gamma and self.normalize: + x = layer_norm(x, self.gamma, self.beta) + elif self.normalize: + x = F.layer_norm(x, (self.head_dim,), self.gamma, self.beta) + elif self.use_gamma and self.use_beta: + x = torch.addcmul(self.beta, x, self.gamma) + elif self.use_gamma: + x = x.mul(self.gamma) + else: + raise RuntimeError(f"Not supported combination of `use_gamma`, `use_beta` and `normalize`, " + f"which is currentlt set as (`{self.use_gamma}`, `{self.use_beta}`, `{self.normalize}`)") + if not flatten: + return x + x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) + # rebased use learnable parameters to approximate any quadratic function + return torch.cat([x2_2 * self.head_dim ** -0.5, x2_1 * (2 / self.head_dim) ** 0.5], dim=-1) + + +class ReLUFeatureMap(nn.Module): + + def __init__( + self, + ) -> ReLUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return F.relu(x) + + +class SquaredReLUFeatureMap(nn.Module): + + def __init__( + self, + ) -> SquaredReLUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return sqrelu(x) + + +class GELUFeatureMap(nn.Module): + + def __init__( + self, + ) -> GELUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return fast_gelu_impl(x) + + +class SwishFeatureMap(nn.Module): + + def __init__( + self, + ) -> SwishFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return swish(x) + + +class SigmoidFeatureMap(nn.Module): + + def __init__( + self, + ) -> SigmoidFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return sigmoid(x) diff --git a/fla/modules/fused_bitlinear.py b/fla/modules/fused_bitlinear.py new file mode 100644 index 0000000000000000000000000000000000000000..c21b7a90d3074567cb69f218481f83d2186599d0 --- /dev/null +++ b/fla/modules/fused_bitlinear.py @@ -0,0 +1,633 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# +# Implementations of BitLinear layer with fused LayerNorm and quantized Linear layer. +# [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764) +# [Scalable MatMul-free Language Modeling](https://arxiv.org/abs/2406.02528) +# +# Code adapted from https://github.com/ridgerchu/matmulfreellm/ + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.modules.layernorm import RMSNorm +from fla.utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard, require_version + +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +def activation_quant(x): + """ + Per-token quantization to 8 bits. No grouping is needed for quantization. + + Args: + x: An activation tensor with shape [n, d]. + + Returns: + A quantized activation tensor with shape [n, d]. + """ + # Compute the scale factor + scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) + # Quantize and then de-quantize the tensor + y = (x * scale).round().clamp_(-128, 127) / scale + return y + + +def weight_quant(w): + """ + Per-tensor quantization to 1.58 bits. No grouping is needed for quantization. + + Args: + w: A weight tensor with shape [d, k]. + + Returns: + A quantized weight tensor with shape [d, k]. + """ + # Compute the scale factor + scale = 1.0 / w.abs().mean().clamp_(min=1e-5) + # Quantize and then de-quantize the tensor + u = (w * scale).round().clamp_(-1, 1) / scale + return u + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_fwd_kernel_quant( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + RESIDUAL, # pointer to the residual + RESIDUAL_OUT, # pointer to the residual + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_res_row, + stride_res_out_row, + N, # number of columns in X + eps, # epsilon to avoid division by zero + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + X += row * stride_x_row + Y += row * stride_y_row + if HAS_RESIDUAL: + RESIDUAL += row * stride_res_row + if STORE_RESIDUAL_OUT: + RESIDUAL_OUT += row * stride_res_out_row + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) + x += residual + if STORE_RESIDUAL_OUT: + tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + + y = x_hat * w if HAS_WEIGHT else x_hat + if HAS_BIAS: + y = y + b + + # Aply quantization to the output + scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5) + # Quantize and then de-quantize the tensor + y = tl.extra.cuda.libdevice.round(y * scale) + y = tl.maximum(tl.minimum(y, 127), -128) / scale + + # Write output + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_fwd_quant( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, +): + if residual is not None: + residual_dtype = residual.dtype + M, N = x.shape + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + residual_out = torch.empty(M, N, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + layer_norm_fwd_kernel_quant[(M,)]( + x, + y, + weight, + bias, + residual, + residual_out, + mean, + rstd, + x.stride(0), + y.stride(0), + residual.stride(0) if residual is not None else 0, + residual_out.stride(0) if residual_out is not None else 0, + N, + eps, + is_rms_norm, + BLOCK_N, + residual is not None, + residual_out is not None, + weight is not None, + bias is not None, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +@triton.heuristics({ + "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DRESIDUAL, + DRESIDUAL_IN, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dres_row, + stride_dres_in_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + rows_per_program, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + row_start = row_block_id * rows_per_program + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += row_start * stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += row_start * stride_dres_in_row + DY += row_start * stride_dy_row + DX += row_start * stride_dx_row + if RECOMPUTE_OUTPUT: + Y += row_start * stride_y_row + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.0) + if RECOMPUTE_OUTPUT: + y = xhat * w if HAS_WEIGHT else xhat + if HAS_BIAS: + y = y + b + + # Aply quantization to the output + scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5) + # Quantize and then de-quantize the tensor + y = tl.extra.cuda.libdevice.round(y * scale) + y = tl.maximum(tl.minimum(y, 127), -128) / scale + + tl.store(Y + cols, y, mask=mask) + wdy = dy + if HAS_WEIGHT: + wdy = dy * w + dw += dy * xhat + if HAS_BIAS: + db += dy + if not IS_RMS_NORM: + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + c1 = tl.sum(xhat * wdy, axis=0) / N + dx = (wdy - xhat * c1) * rstd + if HAS_DRESIDUAL: + dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) + dx += dres + # Write dx + if STORE_DRESIDUAL: + tl.store(DRESIDUAL_IN + cols, dx, mask=mask) + tl.store(DX + cols, dx, mask=mask) + + X += stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += stride_dres_in_row + if RECOMPUTE_OUTPUT: + Y += stride_y_row + DY += stride_dy_row + DX += stride_dx_row + if HAS_WEIGHT: + tl.store(DW + row_block_id * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * N + cols, db, mask=mask) + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + dresidual: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, +): + M, N = x.shape + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device) + dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + sm_count = get_multiprocessor_count(x.device.index) + _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) if weight is not None else None + _db = torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) if bias is not None else None + rows_per_program = math.ceil(M / sm_count) + grid = (sm_count,) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + y, + dy, + dx, + _dw, + _db, + dresidual, + dresidual_in, + mean, + rstd, + x.stride(0), + 0 if not recompute_output else y.stride(0), + dy.stride(0), + dx.stride(0), + dresidual.stride(0) if dresidual is not None else 0, + dresidual_in.stride(0) if dresidual_in is not None else 0, + M, + N, + eps, + rows_per_program, + is_rms_norm, + BLOCK_N, + dresidual is not None, + dresidual_in is not None, + weight is not None, + bias is not None, + ) + dw = _dw.sum(0).to(weight.dtype) if weight is not None else None + db = _db.sum(0).to(bias.dtype) if bias is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype: + dresidual_in = dx + return (dx, dw, db, dresidual_in) if not recompute_output else (dx, dw, db, dresidual_in, y) + + +class LayerNormLinearQuantFn(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + ): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None) + y, mean, rstd, residual_out = layer_norm_fwd_quant( + x, + norm_weight, + norm_bias, + eps, + residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = weight_quant(linear_weight).to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd( + dy, + x, + norm_weight, + norm_bias, + ctx.eps, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y) + return ( + dx.reshape(ctx.x_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +def layer_norm_linear_quant_fn( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, +): + return LayerNormLinearQuantFn.apply( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + ) + + +def rms_norm_linear_quant( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, +): + return layer_norm_linear_quant_fn( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True, + ) + + +@require_version("triton>=3.0", "Triton >= 3.0 is required to do online quantization.") +def bit_linear(x, weight, bias=None, norm_weight=None, norm_bias=None, eps=1e-8): + """ + A functional version of BitLinear that applies quantization to activations and weights. + + Args: + x: Input tensor with shape [n, d]. + weight: Weight tensor with shape [out_features, in_features]. + bias: Bias tensor with shape [out_features] (optional). + norm_weight: Weight tensor for RMS normalization with shape [in_features]. + norm_bias: Bias tensor for RMS normalization with shape [in_features]. + eps: A small constant for numerical stability in normalization. + + Returns: + Output tensor with shape [n, out_features]. + """ + return layer_norm_linear_quant_fn( + x, + norm_weight, + norm_bias, + weight, + bias, + is_rms_norm=True, + ) + + +class BitLinear(nn.Linear): + """ + A custom linear layer that applies quantization on both activations and weights. + This is primarily for training; kernel optimization is needed for efficiency in deployment. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + norm_eps: float = 1e-8, + ): + """ + Initializes the BitLinear layer. + + Args: + in_features: Size of each input sample. + out_features: Size of each output sample. + bias: If set to False, the layer will not learn an additive bias. Default: True. + """ + # Initialize the superclass nn.Linear with the given parameters + super().__init__(in_features, out_features, bias=bias) + + self.norm = RMSNorm(in_features, eps=norm_eps, dtype=torch.float32) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({super().extra_repr()}, norm_eps={self.norm.eps})" + + def forward(self, x): + """ + Overrides the forward pass to include quantization. + + Args: + x: An input tensor with shape [n, d]. + + Returns: + An output tensor with shape [n, d]. + """ + # Weight tensor + w = self.weight + + # Apply RMS normalization to the input + x_norm = self.norm(x) + + # Apply quantization to both activations and weights + # Uses Straight-Through Estimator (STE) trick with .detach() for gradient flow + x_quant = x_norm + (activation_quant(x_norm) - x_norm).detach() + w_quant = w + (weight_quant(w) - w).detach() + # Perform linear operation with quantized values + y = F.linear(x_quant, w_quant) + + return y + + +class FusedBitLinear(BitLinear): + """ + A custom linear layer that applies quantization on both activations and weights. + This is primarily for training; kernel optimization is needed for efficiency in deployment. + """ + + def __init__(self, in_features, out_features, bias=False): + """ + Initializes the BitLinear layer. + + Args: + in_features: Size of each input sample. + out_features: Size of each output sample. + bias: If set to False, the layer will not learn an additive bias. Default: True. + """ + # Initialize the superclass nn.Linear with the given parameters + super().__init__(in_features, out_features, bias=bias) + + def forward(self, x): + return layer_norm_linear_quant_fn( + x, + self.norm.weight, + self.norm.bias, + self.weight, + self.bias, + is_rms_norm=True, + ) diff --git a/fla/modules/fused_cross_entropy.py b/fla/modules/fused_cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..aea0e68bfd110a1988b99580568ffa0981338756 --- /dev/null +++ b/fla/modules/fused_cross_entropy.py @@ -0,0 +1,418 @@ + +# Copyright (c) 2023, Tri Dao. + +from typing import Any + +import torch +import torch.nn as nn +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import input_guard + +# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for +# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent +# version of PyTorch. The following 2 lines are for backward compatibility with +# older PyTorch. +if "all_gather_into_tensor" not in dir(torch.distributed): + torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base + + +@triton.heuristics({ + "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0, +}) +@triton.jit +def cross_entropy_fwd_kernel( + loss_ptr, # data ptrs + lse_ptr, + z_loss_ptr, + logits_ptr, + labels_ptr, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + n_rows, + logits_row_stride, # strides + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, + # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE + SPLIT: tl.constexpr, +): + row_idx = tl.program_id(0) + col_block_idx = tl.program_id(1) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + label_idx = tl.load(labels_ptr + row_idx) + logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")) + logits = logits.to(tl.float32) * logit_scale + max_logits = tl.max(logits, 0) + if HAS_SMOOTHING: + sum_logits = tl.sum(tl.where(col_offsets < n_cols, logits, 0.0), 0) + lse = log(tl.sum(exp(logits - max_logits), 0)) + max_logits + tl.store(lse_ptr + col_block_idx * n_rows + row_idx, lse) + if label_idx == ignore_index: + loss = 0.0 + z_loss = 0.0 + else: + label_idx -= class_start_idx + if label_idx >= col_block_idx * BLOCK_SIZE and label_idx < min( + n_cols, (col_block_idx + 1) * BLOCK_SIZE, + ): + logits_label = tl.load(logits_ptr + label_idx) * logit_scale + if HAS_SMOOTHING: + loss = ( + (lse if not SPLIT else 0.0) + - label_smoothing * sum_logits / total_classes + - (1 - label_smoothing) * logits_label + ) + else: + loss = (lse if not SPLIT else 0.0) - logits_label + else: + # If label is out of bounds, we set the CE loss to 0.0. But we still want the label_smoothing loss + if HAS_SMOOTHING: + loss = label_smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) + else: + loss = 0.0 + if not SPLIT: + z_loss = lse_square_scale * lse * lse + loss += z_loss + else: + z_loss = 0.0 + tl.store(loss_ptr + col_block_idx * n_rows + row_idx, loss) + if not SPLIT: + tl.store(z_loss_ptr + col_block_idx * n_rows + row_idx, z_loss) + + +@triton.heuristics({ + "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0, +}) +@triton.jit +def cross_entropy_bwd_kernel( + dlogits_ptr, # data ptrs + dloss_ptr, + logits_ptr, + lse_ptr, + labels_ptr, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + logits_row_stride, # strides + dlogits_row_stride, + dloss_row_stride, + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, +): + row_idx = tl.program_id(0) + col_block_idx = tl.program_id(1) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) + col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + label_idx = tl.load(labels_ptr + row_idx) + if label_idx != ignore_index: + dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) + else: + dloss = 0.0 + logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( + tl.float32, + ) * logit_scale + lse = tl.load(lse_ptr + row_idx) + probs = exp(logits - lse) + probs += 2.0 * lse_square_scale * lse * probs + label_idx -= class_start_idx + if HAS_SMOOTHING: + smooth_negative = label_smoothing / total_classes + probs = tl.where(col_offsets == label_idx, probs - (1 - label_smoothing), probs) - smooth_negative + else: + probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) + tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) + + +def fused_cross_entropy_forward( + logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + ignore_index: int = -100, + process_group=None, +): + n_rows, n_cols = logits.shape + assert target.shape == (n_rows,) + world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) + total_classes = world_size * n_cols + rank = 0 if process_group is None else torch.distributed.get_rank(process_group) + class_start_idx = rank * n_cols + + if logits.stride(-1) != 1: + logits = logits.contiguous() + # Set these similar to https://github.com/triton-lang/triton/blob/main/python/tutorials/02-fused-softmax.py + MAX_BLOCK_SIZE = 64 * 1024 + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) + num_warps = ( + 4 + if BLOCK_SIZE < 2048 + else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) + ) + # We may split the lse computation across multiple blocks, then do a reduction + # lse(local_lse) to get the final LSE. This is faster for large n_cols (e.g., > 64k) + # where having just one thread block processing more than 64k elements is slow. + split = world_size > 1 or n_cols > MAX_BLOCK_SIZE + n_splits = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE + loss_shape = (n_splits, n_rows) if n_splits > 1 else (n_rows,) + losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + lse = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + z_losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + + cross_entropy_fwd_kernel[(n_rows, n_splits)]( + losses, # data ptrs + lse, + z_losses, + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, + n_cols, # shapes + n_rows, + logits.stride(0), # strides + BLOCK_SIZE=BLOCK_SIZE, # constants + num_warps=num_warps, + SPLIT=split, + ) + + if split: + # If there's no label_smoothing, if target are in the vocab of this partition, losses contains + # - predicted logit, and 0 otherwise. + # If there's label_smoothing=0.1, for target in the vocab of this partition, losses contains + # -0.9 * predicted logit - 0.1 * sum logit / total_classes. + # For target not in the vocab of this partition, losses contains + # -0.1 * sum logit / total_classes. + if n_splits > 1: + lse = torch.logsumexp(lse, dim=0) + losses = losses.sum(dim=0) + if world_size > 1: + lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) + torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) + handle_losses = torch.distributed.all_reduce( + losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True, + ) + lse = torch.logsumexp(lse_allgather, dim=0) + handle_losses.wait() + # After the allreduce, if there's no label_smoothing, the total losses are - predicted_logit, + # we just have to add the (global) lse. + # If there's label_smoothing=0.1, the total losses are + # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. + # Again, we just have to add the (global) lse. + losses += lse + if lse_square_scale != 0.0: + z_losses = lse_square_scale * lse.square() + z_losses.masked_fill_(target == ignore_index, 0.0) + losses += z_losses + else: + z_losses = torch.zeros_like(losses) + losses.masked_fill_(target == ignore_index, 0.0) + + return losses, z_losses, lse, total_classes, class_start_idx + + +class CrossEntropyLossFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + logits, + target, + label_smoothing=0.0, + logit_scale=1.0, + lse_square_scale=0.0, + ignore_index=-100, + inplace_backward=False, + process_group=None, + ): + losses, z_losses, lse, total_classes, class_start_idx = fused_cross_entropy_forward( + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + process_group, + ) + ctx.save_for_backward(logits, lse, target) + ctx.mark_non_differentiable(z_losses) + ctx.label_smoothing = label_smoothing + ctx.logit_scale = logit_scale + ctx.lse_square_scale = lse_square_scale + ctx.ignore_index = ignore_index + ctx.total_classes = total_classes + ctx.class_start_idx = class_start_idx + ctx.inplace_backward = inplace_backward + + return losses, z_losses + + @staticmethod + @input_guard + def backward(ctx, grad_losses, grad_z_losses): + del grad_z_losses # z_losses are only for logging. + + logits, lse, target = ctx.saved_tensors + dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) + n_rows, n_cols = logits.shape + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) + num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) + def grid(META): return (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa + cross_entropy_bwd_kernel[grid]( + dlogits, # data ptrs + grad_losses, + logits, + lse, + target, + ctx.label_smoothing, + ctx.logit_scale, + ctx.lse_square_scale, + ctx.ignore_index, + ctx.total_classes, + ctx.class_start_idx, + n_cols, # shapes + logits.stride(0), # strides + dlogits.stride(0), + grad_losses.stride(0), + BLOCK_SIZE=BLOCK_SIZE, # constants + num_warps=num_warps, + ) + return dlogits, None, None, None, None, None, None, None, None + + +def cross_entropy_loss( + logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + ignore_index=-100, + inplace_backward: bool = False, + process_group=None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Arguments: + logits: [batch, vocab_size] + target: [batch,] + label_smoothing: float + logit_scale: float. + Multiply logits by this scale before calculating the loss. + lse_square_scale: float. + If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. + This is also referred to as "z-loss". + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + inplace_backward: bool. + If True, we do the backward pass in-place by modifying the logits. + This saves memory. + process_group: + if not None, we're doing Tensor Parallel: each process is responsible for + one part of the vocab. The loss will be aggregated across processes. + Returns: + losses: [batch,], float + z_losses: [batch,], float + """ + return CrossEntropyLossFunction.apply( + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + inplace_backward, + process_group, + ) + + +class FusedCrossEntropyLoss(nn.Module): + def __init__( + self, + ignore_index: int = -100, + reduction: str = "mean", + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + inplace_backward: bool = False, + process_group: Any = None, + return_z_loss: bool = False, + ): + """ + Arguments: + ignore_index: int. If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. + This is also referred to as "z-loss". + inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. + This saves memory. + process_group: if not None, we're doing Tensor Parallel: each process is responsible for + one part of the vocab. The loss will be aggregated across processes. + return_z_loss: bool. If True, we return the component of the loss contributed by + the lse_square_scale value. This value is only for logging and does not support + backprop. + """ + super().__init__() + if reduction not in ["mean", "none", "sum"]: + raise NotImplementedError("Only support reduction = 'mean' or 'none' or 'sum'") + self.ignore_index = ignore_index + self.reduction = reduction + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.lse_square_scale = lse_square_scale + self.inplace_backward = inplace_backward + self.process_group = process_group + self.return_z_loss = return_z_loss + + def forward(self, input, target): + """ + Arguments: + input: (batch, vocab_size) + target: (batch,) + Returns: + losses: (batch,) if reduction is 'none', else (1,), dtype float + z_loss: (batch,) if reduction is 'none', else (1,), dtype float (if self.return_z_loss) + """ + assert input.is_cuda and target.is_cuda, "Only support CUDA tensors" + loss, z_loss = cross_entropy_loss( + input, + target, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + lse_square_scale=self.lse_square_scale, + ignore_index=self.ignore_index, + inplace_backward=self.inplace_backward, + process_group=self.process_group, + ) + if self.reduction == "mean": + loss = loss.sum() / (target != self.ignore_index).sum() + elif self.reduction == "sum": + loss = loss.sum() + else: + loss = loss + + if not self.return_z_loss: + return loss + + if self.reduction == "mean": + z_loss = z_loss.sum() / (target != self.ignore_index).sum() + elif self.reduction == "sum": + z_loss = z_loss.sum() + else: + z_loss = z_loss + + return loss, z_loss diff --git a/fla/modules/fused_kl_div.py b/fla/modules/fused_kl_div.py new file mode 100644 index 0000000000000000000000000000000000000000..1a193eb96819b2398844f7b1ba0e6f0fc1885205 --- /dev/null +++ b/fla/modules/fused_kl_div.py @@ -0,0 +1,322 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import IS_AMD, input_guard + +# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 +# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 +# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling +# The optimal maximum block size depends on your hardware, your kernel, and your dtype +MAX_FUSED_SIZE = 65536 // 2 +STATIC_WARPS = 32 if not IS_AMD else 16 + + +@triton.jit +def kl_div_kernel( + logits, + target_logits, + loss, + s_logits, + s_loss, + reduction: tl.constexpr, + N: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr, +): + # https://github.com/triton-lang/triton/issues/1058 + # If N*V is too large, i_n * stride will overflow out of int32, so we convert to int64 + i_n = tl.program_id(0).to(tl.int64) + + logits += i_n * s_logits + target_logits += i_n * s_logits + + # m is the max value. use the notation from the paper + sm = float('-inf') + tm = float('-inf') + # d is the sum. use the notation from the paper + sd, td = 0.0, 0.0 + + NV = tl.cdiv(V, BV) + for iv in range(0, NV): + o_x = iv * BV + tl.arange(0, BV) + # for student + b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf')) + b_sm = tl.max(b_sl) + m_new = tl.maximum(sm, b_sm) + sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new)) + sm = m_new + # for teacher + b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf')) + b_tm = tl.max(b_tl) + m_new = tl.maximum(tm, b_tm) + td = td * exp(tm - m_new) + tl.sum(exp(b_tl - m_new)) + tm = m_new + + b_loss = 0. + # KL(y_true || y) = exp(y_true) * (log(y_true) - log(y)) + for iv in range(0, NV): + o_x = iv * BV + tl.arange(0, BV) + b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf')) + b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf')) + b_sp_log = b_sl - sm - log(sd) + b_tp_log = b_tl - tm - log(td) + b_sp = exp(b_sp_log) + b_tp = exp(b_tp_log) + b_kl = tl.where(o_x < V, b_tp * (b_tp_log - b_sp_log), 0) + b_dl = -b_tp + b_sp + b_loss += tl.sum(b_kl) + if reduction == 'batchmean': + b_dl = b_dl / N + tl.store(logits + o_x, b_dl, mask=o_x < V) + + # Normalize the loss by the number of elements if reduction is 'batchmean' + if reduction == 'batchmean': + b_loss = b_loss / N + + tl.store(loss + i_n * s_loss, b_loss) + + +@triton.jit +def elementwise_mul_kernel( + x, + g, + N: tl.constexpr, + B: tl.constexpr, +): + """ + This function multiplies each element of the tensor pointed by x with the value pointed by g. + The multiplication is performed in-place on the tensor pointed by x. + + Parameters: + x: + Pointer to the input tensor. + g: + Pointer to the gradient output value. + N (int): + The number of columns in the input tensor. + B (int): + The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + i_x = tl.program_id(0).to(tl.int64) + o_x = i_x * B + tl.arange(0, B) + + # Load the gradient output value + b_g = tl.load(g) + b_x = tl.load(x + o_x, mask=o_x < N) + tl.store(x + o_x, b_x * b_g, mask=o_x < N) + + +def fused_kl_div_forward( + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str = 'batchmean', +): + device = x.device + + # ideally, we would like to achieve the same memory consumption as [N, H], + # so the expected chunk size should be: + # NC = ceil(V / H) + # C = ceil(N / NC) + # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048 + N, H, V = *x.shape, weight.shape[0] + BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + # TODO: in real cases, we may need to limit the number of chunks NC to + # ensure the precisions of accumulated gradients + NC = min(8, triton.cdiv(V, H)) + C = triton.next_power_of_2(triton.cdiv(N, NC)) + NC = triton.cdiv(N, C) + + dx = torch.zeros_like(x, device=device) + dw = torch.zeros_like(weight, device=device) if weight is not None else None + # we use fp32 for loss accumulator + loss = torch.zeros(N, dtype=torch.float32, device=device) + + for ic in range(NC): + start, end = ic * C, min((ic + 1) * C, N) + # [C, N] + c_sx = x[start:end] + c_tx = target_x[start:end] + # when doing matmul, use the original precision + # [C, V] + c_sl = F.linear(c_sx, weight) + c_tl = F.linear(c_tx, target_weight) + + # unreduced loss + c_loss = loss[start:end] + + # Here we calculate the gradient of c_sx in place so we can save memory. + kl_div_kernel[(c_sx.shape[0],)]( + logits=c_sl, + target_logits=c_tl, + loss=c_loss, + s_logits=c_sl.stride(-2), + s_loss=c_loss.stride(-1), + reduction=reduction, + N=N, + V=V, + BV=BV, + num_warps=STATIC_WARPS, + ) + + # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V + # thus dx[start: end] should be of shape: C x H + # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only + # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens. + # Thus, we need an additional scaling factor of (n_non_ignore/total) to scale the gradients. + # [C, H] + + dx[start:end] = torch.mm(c_sl, weight) + + if weight is not None: + torch.addmm(input=dw, mat1=c_sl.t(), mat2=c_sx, out=dw) + + loss = loss.sum() + return loss, dx, dw + + +def fused_kl_div_backward( + do: torch.Tensor, + dx: torch.Tensor, + dw: torch.Tensor, +): + # If cross entropy is the last layer, do is 1.0. Skip the mul to save time + if torch.ne(do, torch.tensor(1.0, device=do.device)): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + N, H = dx.shape + B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) + + elementwise_mul_kernel[(triton.cdiv(N * H, B),)]( + x=dx, + g=do, + N=N*H, + B=B, + num_warps=STATIC_WARPS, + ) + + # handle dw + if dw is not None: + V, H = dw.shape + elementwise_mul_kernel[(triton.cdiv(V * H, B),)]( + x=dw, + g=do, + N=V*H, + B=B, + num_warps=STATIC_WARPS, + ) + + return dx, dw + + +class FusedKLDivLossFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str, + ): + loss, dx, dw = fused_kl_div_forward( + x=x, + target_x=target_x, + weight=weight, + target_weight=target_weight, + reduction=reduction, + ) + ctx.save_for_backward(dx, dw) + return loss + + @staticmethod + @input_guard + def backward(ctx, do): + dx, dw = ctx.saved_tensors + dx, dw = fused_kl_div_backward(do, dx, dw) + return dx, None, dw, None, None + + +def fused_kl_div_loss( + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str = 'batchmean', +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target_x (torch.Tensor): [batch_size * seq_len, hidden_size] + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + target_weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + reduction: + Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'. + Returns: + loss + """ + return FusedKLDivLossFunction.apply( + x, + target_x, + weight, + target_weight, + reduction, + ) + + +class FusedKLDivLoss(nn.Module): + + def __init__( + self, + reduction: str = 'batchmean', + ): + """ + Args: + reduction: + Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'. + """ + super().__init__() + + assert reduction in ['batchmean'], f"reduction: {reduction} is not supported" + + self.reduction = reduction + + def forward( + self, + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + ): + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target_x (torch.Tensor): [batch_size * seq_len, hidden_size] + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + target_weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + Returns: + loss + """ + loss = fused_kl_div_loss( + x=x, + target_x=target_x, + weight=weight, + target_weight=target_weight, + reduction=self.reduction, + ) + return loss diff --git a/fla/modules/fused_linear_cross_entropy.py b/fla/modules/fused_linear_cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..f7eed80bdae65f9d2b4fdb8db8b6c90ce35cbeca --- /dev/null +++ b/fla/modules/fused_linear_cross_entropy.py @@ -0,0 +1,619 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Code adapted from +# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py + +from functools import partial + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from torch.distributed import DeviceMesh +from torch.distributed.tensor import Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.ops.utils import logsumexp_fwd +from fla.ops.utils.op import exp +from fla.utils import IS_AMD, input_guard + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + +# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 +# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 +# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling +# The optimal maximum block size depends on your hardware, your kernel, and your dtype +MAX_FUSED_SIZE = 65536 // 2 +STATIC_WARPS = 32 if not IS_AMD else 16 + + +@triton.jit +def cross_entropy_kernel( + logits, + lse, + target, + loss, + total, + ignore_index, + label_smoothing: tl.constexpr, + logit_scale: tl.constexpr, + reduction: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr, +): + """ + This kernel computes both cross entropy loss and the gradient of the input. + We only consider hard label + mean reduction for now. + Please refer to https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html for the math. + + Args: + logits: + Pointer to logits tensor. + lse: + Pointer to logsumexp tensor. + target: Pointer to target tensor. + loss: + Pointer to tensor to store the loss. + V (int): + The number of columns in the input tensor. + total (int): + The number of non-ignored classes. + ignore_index (int): + The index to ignore in the target. + label_smoothing (float): + The amount of smoothing when computing the loss, where 0.0 means no smoothing. + reduction (str): + The string for the reduction to apply + BV (int): + The block size for vocab. + """ + + # https://github.com/triton-lang/triton/issues/1058 + # If B*T*V is too large, i_n * stride will overflow out of int32, so we convert to int64 + i_n = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + + # 1. Load target first because if the target is ignore_index, we can return right away + b_y = tl.load(target + i_n) + + # 2. locate the start index + logits += i_n * V + + if b_y == ignore_index: + # set all x as 0 + for i in range(0, V, BV): + o_v = i + tl.arange(0, BV) + tl.store(logits + o_v, 0.0, mask=o_v < V) + return + + # Online softmax: 2 loads + 1 store (compared with 3 loads + 1 store for the safe softmax) + # Refer to Algorithm 3 in the paper: https://arxiv.org/pdf/1805.02867 + + # 3. [Online softmax] first pass: compute logsumexp + # we did this in anouter kernel + b_l = tl.load(logits + b_y) * logit_scale + b_lse = tl.load(lse + i_n) + + # 4. Calculate the loss + # loss = lse - logits_l + b_loss = b_lse - b_l + + # Label smoothing is a general case of normal cross entropy + # See the full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issue-2503665310 + b_z = 0.0 + eps = label_smoothing / V + + # We need tl.debug_barrier() as mentioned in + # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/ops/cross_entropy.py#L34 + tl.debug_barrier() + + # 5. [Online Softmax] Second pass: compute gradients + # For 'mean' reduction, gradients are normalized by number of non-ignored elements + # dx_y = (softmax(x_y) - 1) / N + # dx_i = softmax(x_i) / N, i != y + # For label smoothing: + # dx_i = (softmax(x_y) - label_smoothing / V) / N, i != y + # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N + # = dx_i - (1 - label_smoothing) / N + for iv in range(0, NV): + o_v = iv * BV + tl.arange(0, BV) + b_logits = tl.load(logits + o_v, mask=o_v < V, other=float('-inf')) * logit_scale + if label_smoothing > 0: + # scale X beforehand to avoid overflow + b_z += tl.sum(tl.where(o_v < V, -eps * b_logits, 0.0)) + b_p = (exp(b_logits - b_lse) - eps) * logit_scale + if reduction == "mean": + b_p = b_p / total + tl.store(logits + o_v, b_p, mask=o_v < V) + + tl.debug_barrier() + + # Orginal loss = H(q, p), with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps + # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p) + # = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i)) + # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as: + # = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd)) + # Refer to H(q', p) in section 7 of the paper: + # https://arxiv.org/pdf/1512.00567 + # pytorch: + # https://github.com/pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L516 + # See full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issuecomment-2333753087 + if label_smoothing > 0: + b_loss = b_loss * (1 - label_smoothing) + (b_z + label_smoothing * b_lse) + + # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N` + b_l = tl.load(logits + b_y) + + # Normalize the loss by the number of non-ignored elements if reduction is "mean" + if reduction == 'mean': + b_loss = b_loss / total + b_l += (label_smoothing - 1) / total * logit_scale + else: + b_l += (label_smoothing - 1) * logit_scale + + tl.store(loss + i_n, b_loss) + tl.store(logits + b_y, b_l) + + +@triton.jit +def elementwise_mul_kernel( + x, + g, + N: tl.constexpr, + B: tl.constexpr, +): + """ + This function multiplies each element of the tensor pointed by x with the value pointed by g. + The multiplication is performed in-place on the tensor pointed by x. + + Parameters: + x: + Pointer to the input tensor. + g: + Pointer to the gradient output value. + N (int): + The number of columns in the input tensor. + B (int): + The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + i_x = tl.program_id(0).to(tl.int64) + o_x = i_x * B + tl.arange(0, B) + + # Load the gradient output value + b_g = tl.load(g) + b_x = tl.load(x + o_x, mask=o_x < N) + tl.store(x + o_x, b_x * b_g, mask=o_x < N) + + +def fused_linear_cross_entropy_forward( + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean", + use_l2warp: bool = False, + l2_penalty_factor: float = 1e-4, +): + device = x.device + # inputs have shape: [N, H] + # materialized activations will have shape: [N, V] + # the increase in memory = [N, V] + # reduction can be achieved by partitioning the number of tokens N into smaller chunks. + + # ideally, we would like to achieve the same memory consumption as [N, H], + # so the expected chunk size should be: + # NC = ceil(V / H) + # C = ceil(N / NC) + # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048 + N, H, V = *x.shape, weight.shape[0] + BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + # TODO: in real cases, we may need to limit the number of chunks NC to + # ensure the precisions of accumulated gradients + NC = min(num_chunks, triton.cdiv(V, H)) + C = triton.next_power_of_2(triton.cdiv(N, NC)) + NC = triton.cdiv(N, C) + + # [N, H] + dx = torch.zeros_like(x, device=device) + # [V, H] + dw = torch.zeros_like(weight, device=device, dtype=torch.float) if weight is not None else None + # [V] + db = torch.zeros_like(bias, device=device, dtype=torch.float) if bias is not None else None + # [N] + loss = torch.zeros(N, device=device, dtype=torch.float) + + total = target.ne(ignore_index).sum().item() + + for ic in range(NC): + start, end = ic * C, min((ic + 1) * C, N) + # [C, N] + c_x = x[start:end] + # when doing matmul, use the original precision + # [C, V] + c_logits = F.linear(c_x, weight, bias) + c_target = target[start:end] + # [C] + # keep lse in fp32 to maintain precision + c_lse = logsumexp_fwd(c_logits, scale=logit_scale, dtype=torch.float) + + # unreduced loss + c_loss = loss[start:end] + if use_l2warp: + c_maxx, c_ids = torch.max(c_logits, -1, keepdim=True) + + # Here we calculate the gradient of c_logits in place so we can save memory. + cross_entropy_kernel[(c_logits.shape[0],)]( + logits=c_logits, + lse=c_lse, + target=c_target, + loss=c_loss, + total=total, + ignore_index=ignore_index, + label_smoothing=label_smoothing, + logit_scale=logit_scale, + reduction=reduction, + V=V, + BV=BV, + num_warps=STATIC_WARPS, + ) + if use_l2warp: + # a. Calculate the L2 gradient w.r.t logits (g_logits_l2) + g_logits_l2 = torch.zeros_like(c_logits) + + # Normalize factor by B*T, which is the 'total' variable here + l2_factor = l2_penalty_factor / total if reduction == 'mean' else l2_penalty_factor + penalty_grad = c_maxx * l2_factor + g_logits_l2.scatter_(-1, c_ids, penalty_grad) + + # b. Backpropagate g_logits_l2 to get its effect on dx, dw, db + # and add it to the main gradients. + # Total_dx = CE_dx + L2_dx + # Total_dw = CE_dw + L2_dw + # Total_db = CE_db + L2_db + if weight is not None: + dw.add_(g_logits_l2.t() @ c_x) + if bias is not None: + db.add_(g_logits_l2.sum(0)) + # The dx contribution must be added to the final dx calculation + dx_l2_contribution = torch.mm(g_logits_l2, weight) + else: + dx_l2_contribution = 0.0 + + # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V + # thus dx should be of shape: C x H + dx[start:end] = torch.mm(c_logits, weight) + dx_l2_contribution + + # keep dw in fp32 to maintain precision + if weight is not None: + dw += c_logits.t() @ c_x + + if bias is not None: + torch.add(input=db, other=c_logits.sum(0), out=db) + + loss = loss.sum() + if dw is not None: + dw = dw.to(weight) + if db is not None: + db = db.to(bias) + return loss, dx, dw, db + + +def fused_linear_cross_entropy_backward( + do: torch.Tensor, + dx: torch.Tensor, + dw: torch.Tensor, + db: torch.Tensor, +): + # If cross entropy is the last layer, do is 1.0. Skip the mul to save time + if torch.ne(do, torch.tensor(1.0, device=do.device)): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + N, H = dx.shape + B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) + + elementwise_mul_kernel[(triton.cdiv(N * H, B),)]( + x=dx, + g=do, + N=N*H, + B=B, + num_warps=STATIC_WARPS, + ) + + # handle dw + if dw is not None: + V, H = dw.shape + elementwise_mul_kernel[(triton.cdiv(V * H, B),)]( + x=dw, + g=do, + N=V*H, + B=B, + num_warps=STATIC_WARPS, + ) + + if db is not None: + V = db.shape[0] + elementwise_mul_kernel[(triton.cdiv(V, B),)]( + x=db, + g=do, + N=V, + B=B, + num_warps=STATIC_WARPS, + ) + return dx, dw, db + + +class FusedLinearCrossEntropyFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean", + use_l2warp: bool = False, + l2_penalty_factor: float = 1e-4, + ): + """ + Fusing the last linear layer with cross-entropy loss + Reference: https://github.com/mgmalek/efficient_cross_entropy + + Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding + the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can + compute the gradient at the forward pass. By doing so, we don't have to store the x and target + for the backward pass. + + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target (torch.LongTensor): [batch_size * seq_len] + where each value is in [0, vocab_size). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + ignore_index: + the index to ignore in the target. + label_smoothing: + the amount of smoothing when computing the loss, where 0.0 means no smoothing. + logit_scale: float = 1.0, + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + use_l2warp: bool = False, + Whether to use L2 regularization on the logits to prevent overconfidence. + Default: False + l2_penalty_factor: float = 1e-4, + """ + loss, dx, dw, db = fused_linear_cross_entropy_forward( + x, + target, + weight, + bias, + ignore_index, + label_smoothing, + logit_scale, + num_chunks, + reduction, + use_l2warp, + l2_penalty_factor, + ) + # downcast to dtype and store for backward + ctx.save_for_backward( + dx.detach(), + dw.detach() if weight is not None else None, + db.detach() if bias is not None else None, + ) + return loss + + @staticmethod + @input_guard + def backward(ctx, do): + dx, dw, db = ctx.saved_tensors + dx, dw, db = fused_linear_cross_entropy_backward(do, dx, dw, db) + return dx, None, dw, db, None, None, None, None, None, None, None + + +def fused_linear_cross_entropy_loss( + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean", + use_l2warp: bool = False, + l2_penalty_factor: float = 1e-4, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target (torch.LongTensor): [batch_size * seq_len] + where each value is in [0, vocab_size). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + Returns: + losses: [batch,], float + """ + return FusedLinearCrossEntropyFunction.apply( + x, + target, + weight, + bias, + ignore_index, + label_smoothing, + logit_scale, + num_chunks, + reduction, + use_l2warp, + l2_penalty_factor, + ) + + +class FusedLinearCrossEntropyLoss(nn.Module): + + def __init__( + self, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean", + use_l2warp: bool = False, + l2_penalty_factor: float = 1e-4, + ): + """ + Args: + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + """ + super().__init__() + + assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" + + self.ignore_index = ignore_index + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.num_chunks = num_chunks + self.reduction = reduction + self.use_l2warp = use_l2warp + self.l2_penalty_factor = l2_penalty_factor + + @torch.compiler.disable + def forward( + self, + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): [batch_size, seq_len, hidden_size] + target (torch.LongTensor): [batch_size, seq_len] + where each value is in [0, V). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + Returns: + loss + """ + loss = fused_linear_cross_entropy_loss( + x.view(-1, x.shape[-1]), + target.view(-1), + weight=weight, + bias=bias, + ignore_index=self.ignore_index, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + num_chunks=self.num_chunks, + reduction=self.reduction, + use_l2warp=self.use_l2warp, + l2_penalty_factor=self.l2_penalty_factor, + ) + return loss + + +class LinearLossParallel(ParallelStyle): + def __init__( + self, + *, + sequence_dim: int = 1, + use_local_output: bool = False, + ): + super().__init__() + + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + x, target, weight, bias = inputs + + if not isinstance(x, DTensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + x = DTensor.from_local(x, device_mesh, sequence_sharding) + if x.placements != sequence_sharding: + x = x.redistribute(placements=sequence_sharding, async_op=True) + if not isinstance(target, DTensor): + target = DTensor.from_local(target, device_mesh, [Replicate()]) + if target.placements != sequence_sharding: + target = target.redistribute(placements=sequence_sharding, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, [Replicate()]) + if weight.placements != [Replicate()]: + # we replicate the weight/bias in FLCE + weight = weight.redistribute(placements=[Replicate()], async_op=True) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, [Replicate()]) + if bias is not None and bias.placements != [Replicate()]: + bias = bias.redistribute(placements=[Replicate()], async_op=True) + + return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias + + @staticmethod + def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.sequence_sharding), + output_fn=partial(self._prepare_output_fn, self.use_local_output), + ) diff --git a/fla/modules/fused_norm_gate.py b/fla/modules/fused_norm_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..87da22b6a41bbe65c699fba97aa6757baf7cfc0d --- /dev/null +++ b/fla/modules/fused_norm_gate.py @@ -0,0 +1,1240 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard + + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for BT in [16, 32, 64] for num_warps in [4, 8, 16]], + key=["D", "NB", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr(residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr(residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None] + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [2, 4, 8, 16]], + key=["D", "IS_RMS_NORM", "STORE_RESIDUAL_OUT", "HAS_RESIDUAL", "HAS_WEIGHT"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + + +@triton.heuristics( + { + "HAS_DRESIDUAL": lambda args: args["dresidual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + "RECOMPUTE_OUTPUT": lambda args: args["y"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for BT in [16, 32, 64] for num_warps in [4, 8, 16]], + key=["D", "NB", "IS_RMS_NORM", "HAS_DRESIDUAL", "HAS_WEIGHT"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_gated_bwd_kernel( + x, # pointer to the input + g, # pointer to the gate + w, # pointer to the weights + b, # pointer to the biases + y, # pointer to the output to be recomputed + dy, # pointer to the output gradient + dx, # pointer to the input gradient + dg, # pointer to the gate gradient + dw, # pointer to the partial sum of weights gradient + db, # pointer to the partial sum of biases gradient + dresidual, + dresidual_in, + mean, + rstd, + T, + BS, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + i_s = tl.program_id(0) + o_d = tl.arange(0, BD) + m_d = o_d < D + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + b_dw = tl.zeros((BT, BD), dtype=tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d, other=0.0).to(tl.float32) + b_db = tl.zeros((BT, BD), dtype=tl.float32) + + # the caller guarantees NS = min(SM, T), so every program has at least one token. + # the last program's range may slightly exceed T (since BS = ceil(T/NS)); + # make_block_ptr uses the true tensor shape (T, D), so boundary_check + # handles the partial tail tile by zero-padding loads and skipping stores. + # the m_t mask below further ensures dw/db only accumulate valid rows (< T). + for i_t in range(i_s * BS, i_s * BS + BS, BT): + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + # [BT, BD] + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + + if not IS_RMS_NORM: + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t,), (BT,), (0,)) + b_mean = tl.load(p_mean, boundary_check=(0,)) + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t,), (BT,), (0,)) + b_rstd = tl.load(p_rstd, boundary_check=(0,)) + # Compute dx + b_xhat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None] + b_xhat = tl.where(m_d[None, :], b_xhat, 0.0) + + b_y = b_xhat * b_w[None, :] if HAS_WEIGHT else b_xhat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + if RECOMPUTE_OUTPUT: + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + b_sigmoid_g = tl.sigmoid(b_g) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g)) + b_dy = b_dy * b_g * b_sigmoid_g + elif ACTIVATION == "sigmoid": + b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g) + b_dy = b_dy * b_sigmoid_g + b_wdy = b_dy + + if HAS_WEIGHT or HAS_BIAS: + # when BT > BS, a tile may span into the next program's range; + # mask to this program's upper bound to avoid double-counting dw/db. + m_t = (i_t + tl.arange(0, BT)) < min(i_s * BS + BS, T) + if HAS_WEIGHT: + b_wdy = b_dy * b_w + b_dw += tl.where(m_t[:, None], b_dy * b_xhat, 0.0) + if HAS_BIAS: + b_db += tl.where(m_t[:, None], b_dy, 0.0) + if not IS_RMS_NORM: + b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D + b_c2 = tl.sum(b_wdy, axis=1) / D + b_dx = (b_wdy - (b_xhat * b_c1[:, None] + b_c2[:, None])) * b_rstd[:, None] + else: + b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D + b_dx = (b_wdy - b_xhat * b_c1[:, None]) * b_rstd[:, None] + if HAS_DRESIDUAL: + p_dres = tl.make_block_ptr(dresidual, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + b_dres = tl.load(p_dres, boundary_check=(0, 1)).to(tl.float32) + b_dx += b_dres + # Write dx + if STORE_DRESIDUAL: + p_dres_in = tl.make_block_ptr(dresidual_in, (T, D), (D, 1), (i_t, 0), (BT, BD), (1, 0)) + tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), boundary_check=(0, 1)) + + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + if HAS_WEIGHT: + tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d) + if HAS_BIAS: + tl.store(db + i_s * D + o_d, tl.sum(b_db, axis=0), mask=m_d) + + +@triton.heuristics( + { + "HAS_DRESIDUAL": lambda args: args["dresidual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + "RECOMPUTE_OUTPUT": lambda args: args["y"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [2, 4, 8, 16]], + key=["D", "IS_RMS_NORM", "STORE_DRESIDUAL", "HAS_DRESIDUAL", "HAS_WEIGHT"], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_gated_bwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + w, # pointer to the weights + b, # pointer to the biases + y, # pointer to the output to be recomputed + dy, # pointer to the output gradient + dx, # pointer to the input gradient + dg, # pointer to the gate gradient + dw, # pointer to the partial sum of weights gradient + db, # pointer to the partial sum of biases gradient + dresidual, + dresidual_in, + mean, + rstd, + T, + BS, + D: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + i_s = tl.program_id(0) + o_d = tl.arange(0, BD) + mask = o_d < D + x += i_s * BS * D + g += i_s * BS * D + if HAS_DRESIDUAL: + dresidual += i_s * BS * D + if STORE_DRESIDUAL: + dresidual_in += i_s * BS * D + dy += i_s * BS * D + dx += i_s * BS * D + dg += i_s * BS * D + if RECOMPUTE_OUTPUT: + y += i_s * BS * D + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=mask).to(tl.float32) + b_dw = tl.zeros((BD,), dtype=tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=mask, other=0.0).to(tl.float32) + b_db = tl.zeros((BD,), dtype=tl.float32) + + for i_t in range(i_s * BS, min(i_s * BS + BS, T)): + # Load data to SRAM + b_x = tl.load(x + o_d, mask=mask, other=0).to(tl.float32) + b_g = tl.load(g + o_d, mask=mask, other=0).to(tl.float32) + b_dy = tl.load(dy + o_d, mask=mask, other=0).to(tl.float32) + + if not IS_RMS_NORM: + b_mean = tl.load(mean + i_t) + b_rstd = tl.load(rstd + i_t) + # Compute dx + b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_xhat = tl.where(mask, b_xhat, 0.0) + + b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat + if HAS_BIAS: + b_y = b_y + b_b + if RECOMPUTE_OUTPUT: + tl.store(y + o_d, b_y, mask=mask) + + b_sigmoid_g = tl.sigmoid(b_g) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g)) + b_dy = b_dy * b_g * b_sigmoid_g + elif ACTIVATION == "sigmoid": + b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g) + b_dy = b_dy * b_sigmoid_g + b_wdy = b_dy + if HAS_WEIGHT: + b_wdy = b_dy * b_w + b_dw += b_dy * b_xhat + if HAS_BIAS: + b_db += b_dy + if not IS_RMS_NORM: + b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D + b_c2 = tl.sum(b_wdy, axis=0) / D + b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd + else: + b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D + b_dx = (b_wdy - b_xhat * b_c1) * b_rstd + if HAS_DRESIDUAL: + b_dres = tl.load(dresidual + o_d, mask=mask, other=0).to(tl.float32) + b_dx += b_dres + # Write dx + if STORE_DRESIDUAL: + tl.store(dresidual_in + o_d, b_dx, mask=mask) + tl.store(dx + o_d, b_dx, mask=mask) + tl.store(dg + o_d, b_dg, mask=mask) + + x += D + g += D + if HAS_DRESIDUAL: + dresidual += D + if STORE_DRESIDUAL: + dresidual_in += D + if RECOMPUTE_OUTPUT: + y += D + dy += D + dx += D + dg += D + if HAS_WEIGHT: + tl.store(dw + i_s * D + o_d, b_dw, mask=mask) + if HAS_BIAS: + tl.store(db + i_s * D + o_d, b_db, mask=mask) + + +def layer_norm_gated_fwd( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, +): + if residual is not None: + residual_dtype = residual.dtype + T, D = x.shape + if residual is not None: + assert residual.shape == (T, D) + if weight is not None: + assert weight.shape == (D,) + if bias is not None: + assert bias.shape == (D,) + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None + rstd = torch.empty((T,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + + if D <= 512: + # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range + # of T before recompiling the kernel. + # NB = triton.cdiv(T, 2048) + NB = triton.cdiv(T, 2048 * 32) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + layer_norm_gated_fwd_kernel[grid]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + T=T, + D=D, + BD=BD, + NB=NB, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + ) + else: + layer_norm_gated_fwd_kernel1[(T,)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + D=D, + BD=BD, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +def layer_norm_gated_bwd( + dy: torch.Tensor, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + eps: float = 1e-5, + mean: torch.Tensor = None, + rstd: torch.Tensor = None, + dresidual: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, +): + T, D = x.shape + assert dy.shape == (T, D) + if dresidual is not None: + assert dresidual.shape == (T, D) + if weight is not None: + assert weight.shape == (D,) + if bias is not None: + assert bias.shape == (D,) + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device) + dg = torch.empty_like(g) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device) + dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # cap program count to T so no program is completely idle. + # without this, high-SM GPUs (e.g. B200, 160 SMs) with small T would + # launch idle programs whose make_block_ptr offsets exceed the tensor shape. + NS = min(get_multiprocessor_count(x.device.index), T) + BS = math.ceil(T / NS) + + dw = torch.empty((NS, D), dtype=torch.float, device=weight.device) if weight is not None else None + db = torch.empty((NS, D), dtype=torch.float, device=bias.device) if bias is not None else None + grid = (NS,) + + if D <= 512: + # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range + # of T before recompiling the kernel. + # NB = triton.cdiv(T, 2048) + NB = triton.cdiv(T, 2048 * 32) + + layer_norm_gated_bwd_kernel[grid]( + x=x, + g=g, + w=weight, + b=bias, + y=y, + dy=dy, + dx=dx, + dg=dg, + dw=dw, + db=db, + dresidual=dresidual, + dresidual_in=dresidual_in, + mean=mean, + rstd=rstd, + T=T, + D=D, + BS=BS, + BD=BD, + NB=NB, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + STORE_DRESIDUAL=dresidual_in is not None, + ) + else: + layer_norm_gated_bwd_kernel1[grid]( + x=x, + g=g, + w=weight, + b=bias, + y=y, + dy=dy, + dx=dx, + dg=dg, + dw=dw, + db=db, + dresidual=dresidual, + dresidual_in=dresidual_in, + mean=mean, + rstd=rstd, + T=T, + D=D, + BS=BS, + BD=BD, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + STORE_DRESIDUAL=dresidual_in is not None, + ) + dw = dw.sum(0).to(weight.dtype) if weight is not None else None + db = db.sum(0).to(bias.dtype) if bias is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype: + dresidual_in = dx + return (dx, dg, dw, db, dresidual_in) if not recompute_output else (dx, dg, dw, db, dresidual_in, y) + + +class LayerNormGatedFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str, + residual: torch.Tensor | None = None, + eps: float = 1e-6, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + ): + x_shape_og = x.shape + g_shape_og = g.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + g = g.reshape(-1, g.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = residual.dtype if residual is not None else (torch.float if residual_in_fp32 else None) + y, mean, rstd, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=weight, + bias=bias, + activation=activation, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + ctx.save_for_backward(residual_out, g, weight, bias, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.g_shape_og = g_shape_og + ctx.activation = activation + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dy, *args): + x, g, weight, bias, mean, rstd = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dg, dw, db, dres_in = layer_norm_gated_bwd( + dy=dy, + x=x, + g=g, + weight=weight, + bias=bias, + activation=ctx.activation, + eps=ctx.eps, + mean=mean, + rstd=rstd, + dresidual=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + ) + return ( + dx.reshape(ctx.x_shape_og), + dg.reshape(ctx.g_shape_og), + dw, + db, + None, + dres_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +class LayerNormGatedLinearFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor | None = None, + eps: float = 1e-6, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + ): + x_shape_og = x.shape + g_shape_og = g.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + g = g.reshape(-1, g.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = residual.dtype if residual is not None else (torch.float if residual_in_fp32 else None) + y, mean, rstd, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=norm_weight, + bias=norm_bias, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = linear_weight.to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, g, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.g_shape_og = g_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, g, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dg, dnorm_weight, dnorm_bias, dres_in, y = layer_norm_gated_bwd( + dy=dy, + x=x, + g=g, + weight=norm_weight, + bias=norm_bias, + eps=ctx.eps, + mean=mean, + rstd=rstd, + dresidual=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y) + return ( + dx.reshape(ctx.x_shape_og), + dg.reshape(ctx.g_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dres_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +def layer_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + return LayerNormGatedFunction.apply( + x, + g, + weight, + bias, + activation, + residual, + eps, + prenorm, + residual_in_fp32, + False, + ) + + +def rms_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + return LayerNormGatedFunction.apply( + x, + g, + weight, + bias, + activation, + residual, + eps, + prenorm, + residual_in_fp32, + True, + ) + + +def layer_norm_swish_gate_linear( + x: torch.Tensor, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + return LayerNormGatedLinearFunction.apply( + x, + g, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + False, + ) + + +def rms_norm_swish_gate_linear( + x, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + return LayerNormGatedLinearFunction.apply( + x, + g, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + True, + ) + + +class FusedLayerNormGated(nn.Module): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + activation: str = "swish", + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedLayerNormGated: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ["swish", "silu", "sigmoid"]: + raise ValueError(f"Unsupported activation: {self.activation}") + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += f", activation={self.activation}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return layer_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class FusedRMSNormGated(nn.Module): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + activation: str = "swish", + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedRMSNormGated: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ["swish", "silu", "sigmoid"]: + raise ValueError(f"Unsupported activation: {self.activation}") + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += f", activation={self.activation}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class FusedLayerNormSwishGate(FusedLayerNormGated): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedLayerNormSwishGate: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + bias=bias, + eps=eps, + device=device, + dtype=dtype, + ) + + +class FusedRMSNormSwishGate(FusedRMSNormGated): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedRMSNormSwishGate: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype, + ) + + +class FusedLayerNormGatedLinear(nn.Module): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedLayerNormGatedLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return layer_norm_swish_gate_linear( + x, + g, + self.weight, + self.bias, + weight, + bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class FusedLayerNormSwishGateLinear(FusedLayerNormGatedLinear): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedLayerNormSwishGateLinear: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype, + ) + + +class FusedRMSNormGatedLinear(nn.Module): + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedRMSNormGatedLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return rms_norm_swish_gate_linear( + x, + g, + self.weight, + self.bias, + weight, + bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class FusedRMSNormSwishGateLinear(FusedRMSNormGatedLinear): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> FusedRMSNormSwishGateLinear: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype, + ) diff --git a/fla/modules/grpo.py b/fla/modules/grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..aa7ea718fbab1e4d47c34f0d82442ccfba55687e --- /dev/null +++ b/fla/modules/grpo.py @@ -0,0 +1,414 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# +# modified from https://github.com/mdy666/mdy_triton/blob/e0a856347bd988e05e0152332bba35f1d33c5b1f/others/grpo/grpo_loss.ipynb +# XHS ID: blueeeee + +# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py +""" +# Get the per-token log probabilities for the completions for the model and the reference model + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits + logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + + input_ids = input_ids[:, -logits_to_keep:] + # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. + # See https://github.com/huggingface/trl/issues/2770 + logits = logits[:, -logits_to_keep:] + return selective_log_softmax(logits, input_ids) # compute logprobs for the input tokens + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + # Compute the per-token log probabilities for the model + + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + + # x - x.detach() allows for preserving gradients from x + advantages = inputs["advantages"] + per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + per_token_loss = -(per_token_loss - self.beta * per_token_kl) + loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + + # Log the metrics + completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() + self._metrics["completion_length"].append(completion_length) + + mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + + return loss +""" + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import IS_AMD, autotune_cache_kwargs, input_guard + +NUM_WARPS_AUTOTUNE = [4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': BLOCK_SIZE}, num_warps=NUM_WARPS, num_stages=NUM_STAGES) + for BLOCK_SIZE in [1024, 2048, 4096, 8192] + for NUM_WARPS in NUM_WARPS_AUTOTUNE + for NUM_STAGES in [1, 2, 4] + ], + key=['B', 'N'], + **autotune_cache_kwargs, +) +@triton.jit +def grpo_fwd_kernel( + logits_ptr, + ref_logp_ptr, + input_ids_ptr, + advantages_ptr, + completion_mask_ptr, + loss_ptr, + lse_ptr, + beta, + save_kl: tl.constexpr, + B, + M, + N, + L, + start_idx, + BLOCK_SIZE: tl.constexpr, +): + row_idx = tl.program_id(0) + + off_b = row_idx // L + N = tl.cast(N, tl.int64) + + loss_ptr += row_idx + + completion_mask_ptr += row_idx + not_skip = tl.load(completion_mask_ptr).to(tl.int1) + if not_skip == 1: + ref_logp_ptr += row_idx + lse_ptr += row_idx + advantages_ptr += off_b + logits_ptr += N * (row_idx + off_b) + input_ids_ptr += row_idx + (off_b+1) * start_idx + base_cols = tl.arange(0, BLOCK_SIZE) + + m_i = -float("inf") + l_i = 0.0 + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + logits = tl.load(logits_ptr+cols, mask=mask, other=-float('inf')).to(tl.float32) + m_ij = tl.max(logits) + new_m_i = tl.maximum(m_i, m_ij) + l_i = l_i * exp(m_i - new_m_i) + tl.sum(exp(logits - new_m_i)) + m_i = new_m_i + lse = log(l_i) + m_i + + idx = tl.load(input_ids_ptr) + x = tl.load(logits_ptr+idx).to(tl.float32) + advantage = tl.load(advantages_ptr).to(tl.float32) + ref_logp = tl.load(ref_logp_ptr) + logp = x - lse + diff = ref_logp - logp + kl = exp(diff) - diff - 1 + loss = kl * beta - advantage + + tl.store(loss_ptr, loss.to(loss_ptr.dtype.element_ty)) + tl.store(lse_ptr, lse.to(lse_ptr.dtype.element_ty)) + if save_kl: + tl.store(loss_ptr+M, kl.to(loss_ptr.dtype.element_ty)) + else: + # store 0 + tl.store(loss_ptr, 0.0) + if save_kl: + tl.store(loss_ptr+M, 0.0) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=NUM_WARPS, num_stages=NUM_STAGES) + for NUM_WARPS in [32] + for NUM_STAGES in [4] + ], + key=['B', 'N'], + **autotune_cache_kwargs, +) +@triton.jit +def grpo_bwd_kernel( + dloss_ptr, + dlogits_ptr, + logits_ptr, + ref_logp_ptr, + input_ids_ptr, + advantages_ptr, + completion_mask_ptr, + lse_ptr, + beta, + B, + N, + L, + start_idx, + BLOCK_SIZE: tl.constexpr, +): + + row_idx = tl.program_id(0) # B*L + off_b = row_idx // L + + N = tl.cast(N, tl.int64) + + dlogits_ptr += N * (row_idx + off_b) + base_cols = tl.arange(0, BLOCK_SIZE) + completion_mask_ptr += row_idx + not_skip = tl.load(completion_mask_ptr).to(tl.int1) + + if not_skip == 1: + lse_ptr += row_idx + dloss_ptr += row_idx + advantages_ptr += off_b + ref_logp_ptr += row_idx + logits_ptr += N * (row_idx + off_b) + input_ids_ptr += row_idx + (off_b+1) * start_idx + dloss = tl.load(dloss_ptr).to(tl.float32) + lse = tl.load(lse_ptr).to(tl.float32) + idx = tl.load(input_ids_ptr) + x = tl.load(logits_ptr+idx).to(tl.float32) + advantage = tl.load(advantages_ptr).to(tl.float32) + ref_logp = tl.load(ref_logp_ptr) + # Need for in-place grad. + tl.debug_barrier() + logp = x - lse + + dlogp = (beta * (-1.0 * exp(ref_logp - logp) + 1) + - advantage) * dloss + + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + logits = tl.load(logits_ptr+cols, mask=mask, other=-float('inf')).to(tl.float32) + probs = exp(logits - lse) + dlogits = tl.where(cols == idx, 1-probs, -probs) * dlogp + + tl.store(dlogits_ptr+cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask) + else: + dlogits = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + + tl.store(dlogits_ptr+cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask) + + +class GrpoLoss(torch.autograd.Function): + + @input_guard + @staticmethod + def forward(ctx, logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl, inplace=True): + ctx.input_shape = logits.shape + B, L_ADD_1, N = ctx.input_shape + L = L_ADD_1 - 1 + M = B * L + input_ids_start_index = input_ids.size(1) - L + + if not save_kl: + loss = torch.empty(B, L, device=logits.device, dtype=torch.float32) + else: + loss = torch.empty(B*2, L, device=logits.device, dtype=torch.float32) + + lse = torch.empty(B, L, device=logits.device, dtype=torch.float32) + + if completion_mask is None: + completion_mask = torch.ones(B, L, device=logits.device, dtype=torch.int32) + else: + loss[:B].masked_fill_(completion_mask.logical_not(), 0.0) + + grpo_fwd_kernel[(M,)]( + logits_ptr=logits, + ref_logp_ptr=ref_logp, + input_ids_ptr=input_ids, + advantages_ptr=advantages, + completion_mask_ptr=completion_mask, + loss_ptr=loss, + lse_ptr=lse, + beta=beta, + save_kl=save_kl, + B=B, M=M, N=N, L=L, + start_idx=input_ids_start_index, + ) + ctx.beta = beta + ctx.save_for_backward(lse, logits, input_ids, advantages, completion_mask) + ctx.ref_logp = ref_logp + ctx.inplace = inplace + return loss + + @input_guard + @staticmethod + def backward(ctx, dloss): + # The grad of logits comes from two parts, the reward part and the kl part + lse, logits, input_ids, advantages, completion_mask = ctx.saved_tensors + inplace = ctx.inplace + B, L_ADD_1, N = ctx.input_shape + L = L_ADD_1 - 1 + M = B * L + + input_ids_start_index = input_ids.size(1) - L + + # B, L_ADD_1, N + dlogits = logits if inplace else torch.empty_like(logits) + BN = min(65536, triton.next_power_of_2(N)) + + grpo_bwd_kernel[(M,)]( + dloss_ptr=dloss, + dlogits_ptr=dlogits, + logits_ptr=logits, + ref_logp_ptr=ctx.ref_logp, + input_ids_ptr=input_ids, + advantages_ptr=advantages, + completion_mask_ptr=completion_mask, + lse_ptr=lse, + beta=ctx.beta, + B=B, N=N, L=L, + BLOCK_SIZE=BN, + start_idx=input_ids_start_index, + ) + # The last token in the completion is not used in the loss computation + # and therefore its gradient should be set to 0 + dlogits[:, -1, :].fill_(0.0) + return dlogits.view(*ctx.input_shape), None, None, None, None, None, None, None + + +def fused_grpo_loss(logits, ref_logp, input_ids, advantages, + beta=0.1, completion_mask=None, save_kl=False, inplace=False) -> torch.Tensor: + ''' + compute grpo loss, save memory(no addition usage) and fast speed(6X for A800) + + Args: + logtits: Tensor, [B, L+1, vocab_size], the origin output of model, it's not logits[:, :-1] + ref_logp: Tensor, [B, L], the origin output of model, it's not ref_logits[:, :-1] + input_ids: Tensor, [B, K+L], it's prompt_completion_id, it contains the prompt ids and output ids + advantages: Tensor, [B], the advantages of each prompt + beta: float, the weight of kl loss + completion_mask: Tensor, loss mask + save_kl: bool, if true will save kl + + Retutn: + loss: Tensor, [B, L], the loss of grpo, it contains the advantage part and kl part + + NOTE: logits(ref_logits) is computed by these steps + logits_to_keep = completion_ids.size(1) + + def get_per_token_logits(model, input_ids, attention_mask, logits_to_keep): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model( + input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1 + ).logits + return logits + + logits = get_per_token_logits(model, prompt_completion_ids, attention_mask, logits_to_keep) + ''' + out = GrpoLoss.apply(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl, inplace) + if not save_kl: + return out + else: + return out.chunk(2, axis=0) + + +def grpo_loss_torch(logits, ref_logp, input_ids, advantages, beta=0.1, completion_mask=None, save_kl=False): + def get_log_probs(logits, input_ids): + per_token_logps = [] + for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):], strict=False): + log_probs = logits_row.log_softmax(dim=-1) + token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) + per_token_logps.append(token_log_prob) + return torch.stack(per_token_logps) + + logits = logits[:, :-1] + per_token_logps = get_log_probs(logits, input_ids) + ref_per_token_logps = ref_logp + per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + + per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + per_token_loss = -(per_token_loss - beta * per_token_kl) + if completion_mask is not None: + per_token_loss *= completion_mask + if save_kl: + per_token_kl *= completion_mask + return per_token_loss if not save_kl else (per_token_loss, per_token_kl) + + +@torch.compile(fullgraph=True) +def grpo_loss_with_old_logps( + logps: torch.Tensor, + ref_logps: torch.Tensor, + old_logps: torch.Tensor, + pad_mask: torch.Tensor, + logits_to_keep: int, + rewards: torch.Tensor, + beta: float = 0.2, + epsilon: float = 0.2, +): + """ + Compute the GRPO (Group Relative Policy Optimization) loss. + + Args: + logps (torch.Tensor): [Batch, Token_length] Log probabilities of the current policy. + ref_logps (torch.Tensor):[Batch, Token_length] Log probabilities of the reference policy. + old_logps (torch.Tensor): [Batch, Token_length] Log probabilities of the old policy. + completion_ids (torch.Tensor): [Batch, Token_length] Completion token IDs (bool). + pad_token_id: Pad token ID. + logits_to_keep (int): Number of logits to keep for masking. + rewards (torch.Tensor): [Batch] Rewards for each generation. + beta (float) = 0.2: A hyperparameter for weighting the KL divergence term. + epsilon (float) = 0.2: An float hyperparameter for clipping the importance weights. + + Returns: + torch.Tensor: The computed GRPO loss. + """ + B = logps.shape[0] + assert B > 1, "Batch * Num generations should be greater than 1" + + rewards_shaped = rewards.view(-1, B) # B,num_generations + advantages = (rewards_shaped - rewards_shaped.mean(dim=1, keepdim=True)) / \ + (rewards_shaped.std(dim=1, keepdim=True) + 1e-8) + advantages = advantages.view(-1) # B*num_generations + # Calculate the per - token KL divergence + per_token_kl = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1 + + # Calculate the ratio of probabilities (importance weights) + # Importance weights are calculated as exp(log_pi_theta - log_pi_theta_old) + importance_weights = torch.exp(logps - old_logps) + + # Clip the importance weights to the range [1 - epsilon, 1 + epsilon] + importance_weights_clipped = torch.clamp(importance_weights, 1 - epsilon, 1 + epsilon) + + # Create a completion mask. It checks which positions are valid based on logits_to_keep + completion_mask = torch.arange(logits_to_keep, device=logps.device)[None, :] >= 0 + + # Combine the completion mask and padding mask + completion_mask = completion_mask & pad_mask # Ensure matching shape + + # Add an extra dimension to advantages to match the shape for element - wise multiplication + advantages = advantages.unsqueeze(1) + + # Calculate the per - token loss. It takes the minimum of the unclipped and clipped importance weights + # and subtracts the KL divergence term weighted by beta, then multiplies by the completion mask + token_loss = -(torch.min(advantages * importance_weights, advantages * + importance_weights_clipped) - beta * per_token_kl) * completion_mask + + # Calculate the final loss by summing the token losses and normalizing by the number of valid tokens + loss = -token_loss.sum() / completion_mask.sum() + + return loss diff --git a/fla/modules/l2norm.py b/fla/modules/l2norm.py new file mode 100644 index 0000000000000000000000000000000000000000..06f4a45df888a404141dbb651a778ce4e74fcb72 --- /dev/null +++ b/fla/modules/l2norm.py @@ -0,0 +1,282 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn as nn +import triton +import triton.language as tl + +from fla.utils import IS_AMD, autotune_cache_kwargs, input_guard + +BT_LIST = [8, 16, 32, 64, 128] +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], + key=["D"], + **autotune_cache_kwargs, +) +@triton.jit +def l2norm_fwd_kernel1( + x, + y, + rstd, + eps, + D, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + # Compute mean and variance + cols = tl.arange(0, BD) + mask = cols < D + + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + tl.store(rstd + i_t, b_rstd) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], + key=["D"], + **autotune_cache_kwargs, +) +@triton.jit +def l2norm_bwd_kernel1( + y, + rstd, + dy, + dx, + eps, + D, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + y += i_t * D + dx += i_t * D + dy += i_t * D + + cols = tl.arange(0, BD) + mask = cols < D + b_y = tl.load(y + cols, mask=mask, other=0.0).to(tl.float32) + b_rstd = tl.load(rstd + i_t).to(tl.float32) + b_dy = tl.load(dy + cols, mask=mask, other=0.0).to(tl.float32) + b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd + tl.store(dx + cols, b_dx, mask=mask) + + +@triton.autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], + key=["D", "NB"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def l2norm_fwd_kernel( + x, + y, + rstd, + eps, + T, + D: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + BT: tl.constexpr, +): + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) + b_y = b_x * b_rstd[:, None] + + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + +@triton.autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], + key=["D", "NB"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def l2norm_bwd_kernel( + y, + rstd, + dy, + dx, + eps, + T, + D: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + BT: tl.constexpr, +): + i_t = tl.program_id(0) + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + + b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) + b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) + b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + + +def l2norm_fwd( + x: torch.Tensor, + eps: float = 1e-6, + output_dtype: torch.dtype | None = None, +): + x_shape_og = x.shape + x = x.view(-1, x.shape[-1]) + # allocate output + if output_dtype is None: + y = torch.empty_like(x) + else: + y = torch.empty_like(x, dtype=output_dtype) + assert y.stride(-1) == 1 + T, D = x.shape[0], x.shape[-1] + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + + rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + if D <= 512: + # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range + # of T before recompiling the kernel. + # NB = triton.cdiv(T, 2048) + NB = triton.cdiv(T, 2048 * 32) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_fwd_kernel[grid]( + x=x, + y=y, + rstd=rstd, + eps=eps, + T=T, + D=D, + BD=BD, + NB=NB, + ) + else: + l2norm_fwd_kernel1[(T,)]( + x=x, + y=y, + rstd=rstd, + eps=eps, + D=D, + BD=BD, + ) + return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) + + +def l2norm_bwd( + y: torch.Tensor, + rstd: torch.Tensor, + dy: torch.Tensor, + eps: float = 1e-6, +): + y_shape_og = y.shape + y = y.view(-1, dy.shape[-1]) + dy = dy.view(-1, dy.shape[-1]) + assert dy.shape == y.shape + # allocate output + dx = torch.empty_like(y) + T, D = y.shape[0], y.shape[-1] + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // y.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + + if D <= 512: + # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range + # of T before recompiling the kernel. + # NB = triton.cdiv(T, 2048) + NB = triton.cdiv(T, 2048 * 32) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_bwd_kernel[grid]( + y=y, + rstd=rstd, + dy=dy, + dx=dx, + eps=eps, + T=T, + D=D, + BD=BD, + NB=NB, + ) + else: + l2norm_bwd_kernel1[(T,)]( + y=y, + rstd=rstd, + dy=dy, + dx=dx, + eps=eps, + D=D, + BD=BD, + ) + + return dx.view(y_shape_og) + + +class L2NormFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward( + ctx, + x, + eps=1e-6, + output_dtype=None, + ): + y, rstd = l2norm_fwd(x, eps, output_dtype) + ctx.eps = eps + ctx.x_dtype = x.dtype + ctx.save_for_backward(y, rstd) + return y + + @staticmethod + @input_guard + def backward(ctx, dy): + y, rstd = ctx.saved_tensors + dx = l2norm_bwd(y, rstd, dy, ctx.eps) + return dx, None, None + + +def l2norm( + x: torch.Tensor, + eps: float = 1e-6, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + return L2NormFunction.apply(x, eps, output_dtype) + + +l2_norm = l2norm + + +class L2Norm(nn.Module): + def __init__( + self, + eps: float = 1e-6, + output_dtype: torch.dtype | None = None, + ): + super().__init__() + self.eps = eps + self.output_dtype = output_dtype + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return l2norm(x, self.eps, self.output_dtype) diff --git a/fla/modules/l2warp.py b/fla/modules/l2warp.py new file mode 100644 index 0000000000000000000000000000000000000000..4cfb432745e68b5af9577e180776cb88454551ce --- /dev/null +++ b/fla/modules/l2warp.py @@ -0,0 +1,37 @@ + +import torch + + +class L2Wrap(torch.autograd.Function): + r""" + This class of penalty prevents the model from becoming overconfident, + thereby mitigating precision loss in BF16. + + This version is memory-optimized by not storing the full logits tensor. + """ + @staticmethod + def forward(ctx, loss, logits, l2_penalty_factor=1e-4): + """ + Forward pass for L2 penalty. + Args: + loss (torch.Tensor): The loss tensor. + logits (torch.Tensor): Shape[B, T, V] The logits tensor. + l2_penalty_factor (float): The factor for L2 penalty. + """ + maxx, ids = torch.max(logits, dim=-1, keepdim=True) + ctx.logits_shape = logits.shape + factor = l2_penalty_factor / (logits.shape[0] * logits.shape[1]) + maxx = maxx * factor + ctx.save_for_backward(maxx, ids) + return loss + + @staticmethod + def backward(ctx, grad_output): + maxx, ids = ctx.saved_tensors + glogits = torch.zeros(ctx.logits_shape, device=grad_output.device, + dtype=grad_output.dtype) + glogits.scatter_(-1, ids, maxx) + return grad_output, glogits, None + + +l2_warp = L2Wrap.apply diff --git a/fla/modules/layernorm.py b/fla/modules/layernorm.py new file mode 100644 index 0000000000000000000000000000000000000000..aa9b9f8de6dbb0b99b1fc0920a518cf0852eaa1f --- /dev/null +++ b/fla/modules/layernorm.py @@ -0,0 +1,1464 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang + +# Copyright (c) 2023, Tri Dao +# https://github.com/state-spaces/mamba/blob/fb7b5310fa865dbd62aa059b1e26f2b431363e2a/mamba_ssm/ops/triton/layernorm.py +# Implement residual + layer_norm / rms_norm. + +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +from __future__ import annotations + +from functools import partial + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange +from torch.distributed import DeviceMesh +from torch.distributed.tensor import Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + + +def layer_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( + dtype, + ) + return out if not prenorm else (out, x) + + +def rms_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = out.to(dtype) + return out if not prenorm else (out, x) + + +def group_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + residual: torch.Tensor = None, + eps: float = 1e-5, + is_rms_norm: bool = False, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + residual = x + x, weight = [ + rearrange(data, "... (g d) -> ... g d", g=num_groups) for data in (x, weight) + ] + if bias is not None: + bias = rearrange(bias, '... (g d) -> ... g d', g=num_groups) + if not is_rms_norm: + mean = x.mean(dim=-1, keepdim=True) + x = x - mean + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = rearrange(out, "... g d -> ... (g d)") + out = out.to(dtype) + return out if not prenorm else (out, residual) + + +class GroupNormRef(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False, + ) -> GroupNormRef: + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False): + return group_norm_ref( + x, + self.weight, + self.bias, + num_groups=self.num_groups, + residual=residual, + eps=self.eps, + is_rms_norm=self.is_rms_norm, + prenorm=prenorm, + upcast=True, + ) + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for BT in [32, 64, 128] + for num_warps in [2, 4, 8] + ], + key=['D', 'NB', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM'], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_fwd_kernel( + x, # pointer to the input + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + res, # pointer to the res + res_out, # pointer to the res + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, + G: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + + o_t = i_t * BT + tl.arange(0, BT) + o_g = o_t % G + o_d = tl.arange(0, BD) + m_d = o_d < D + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr(res, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr(res_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_g[:, None] * D + o_d[None, :], mask=m_d[None, :]).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_g[:, None] * D + o_d[None, :], mask=m_d[None, :]).to(tl.float32) + b_x_hat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None] + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8, 16] + ], + key=['D', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM'], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_fwd_kernel1( + x, # pointer to the input + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + res, # pointer to the res + res_out, # pointer to the res + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + G: tl.constexpr, + D: tl.constexpr, + BD: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + i_g = i_t % G + + x += i_t * D + y += i_t * D + if HAS_RESIDUAL: + res += i_t * D + if STORE_RESIDUAL_OUT: + res_out += i_t * D + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(res + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(res_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + i_g * D + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + i_g * D + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + + +@triton.heuristics({ + 'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for BT in [32, 64] + for num_warps in [2, 4, 8] + ], + key=['D', 'NB', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_bwd_kernel( + x, # pointer to the input + w, # pointer to the weights + b, # pointer to the biases + y, # pointer to the output to be recomputed + dy, # pointer to the output gradient + dx, # pointer to the input gradient + dw, # pointer to the partial sum of weights gradient + db, # pointer to the partial sum of biases gradient + dres, + dres_in, + mean, + rstd, + T, + G: tl.constexpr, + D: tl.constexpr, + BS: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + NB: tl.constexpr, + GS: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + i_s = tl.program_id(0) + i_g, i_sg = i_s // GS, i_s % GS + + o_d = tl.arange(0, BD) + m_d = o_d < D + if HAS_WEIGHT: + b_w = tl.load(w + i_g * D + o_d, mask=m_d).to(tl.float32) + b_dw = tl.zeros((BT, BD), dtype=tl.float32) + if HAS_BIAS: + b_b = tl.load(b + i_g * D + o_d, mask=m_d, other=0.0).to(tl.float32) + b_db = tl.zeros((BT, BD), dtype=tl.float32) + + # Tg: number of tokens per group, used as the logical shape for make_block_ptr. + # for mean/rstd with shape (T,) and stride (G,), the strided view has Tg elements per group. + # the caller guarantees NS capped so every program has work. + # the last program's range may slightly exceed Tg (since BS = cdiv(T, NS)); + # boundary_check handles the partial tail tile, m_t < Tg masks dw/db accumulation. + Tg = T // G + for i_t in range(i_sg * BS, i_sg * BS + BS, BT): + p_x = tl.make_block_ptr(x + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_dy = tl.make_block_ptr(dy + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + p_dx = tl.make_block_ptr(dx + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + # [BT, BD] + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + + if not IS_RMS_NORM: + p_mean = tl.make_block_ptr(mean + i_g, (Tg,), (G,), (i_t,), (BT,), (0,)) + b_mean = tl.load(p_mean, boundary_check=(0,)) + p_rstd = tl.make_block_ptr(rstd + i_g, (Tg,), (G,), (i_t,), (BT,), (0,)) + b_rstd = tl.load(p_rstd, boundary_check=(0,)) + # Compute dx + b_xhat = (b_x - b_mean[:, None]) * b_rstd[:, None] if not IS_RMS_NORM else b_x * b_rstd[:, None] + b_xhat = tl.where(m_d[None, :], b_xhat, 0.0) + + b_y = b_xhat * b_w[None, :] if HAS_WEIGHT else b_xhat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + if RECOMPUTE_OUTPUT: + p_y = tl.make_block_ptr(y + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + b_wdy = b_dy + + if HAS_WEIGHT or HAS_BIAS: + # when BT > BS, a tile may span into the next program's range; + # mask to this program's upper bound to avoid double-counting dw/db. + m_t = (i_t + tl.arange(0, BT)) < min(i_sg * BS + BS, Tg) + if HAS_WEIGHT: + b_wdy = b_dy * b_w + b_dw += tl.where(m_t[:, None], b_dy * b_xhat, 0.0) + if HAS_BIAS: + b_db += tl.where(m_t[:, None], b_dy, 0.0) + if not IS_RMS_NORM: + b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D + b_c2 = tl.sum(b_wdy, axis=1) / D + b_dx = (b_wdy - (b_xhat * b_c1[:, None] + b_c2[:, None])) * b_rstd[:, None] + else: + b_c1 = tl.sum(b_xhat * b_wdy, axis=1) / D + b_dx = (b_wdy - b_xhat * b_c1[:, None]) * b_rstd[:, None] + if HAS_DRESIDUAL: + p_dres = tl.make_block_ptr(dres + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + b_dres = tl.load(p_dres, boundary_check=(0, 1)).to(tl.float32) + b_dx += b_dres + # Write dx + if STORE_DRESIDUAL: + p_dres_in = tl.make_block_ptr(dres_in + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), boundary_check=(0, 1)) + + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + + if HAS_WEIGHT: + tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d) + if HAS_BIAS: + tl.store(db + i_s * D + o_d, tl.sum(b_db, axis=0), mask=m_d) + + +@triton.heuristics({ + 'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8] + ], + key=['D', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_bwd_kernel1( + x, # pointer to the input + w, # pointer to the weights + b, # pointer to the biases + y, # pointer to the output to be recomputed + dy, # pointer to the output gradient + dx, # pointer to the input gradient + dw, # pointer to the partial sum of weights gradient + db, # pointer to the partial sum of biases gradient + dres, + dres_in, + mean, + rstd, + T, + G: tl.constexpr, + D: tl.constexpr, + BS: tl.constexpr, + BD: tl.constexpr, + GS: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + i_s = tl.program_id(0) + i_g, i_sg = i_s // GS, i_s % GS + + o_d = tl.arange(0, BD) + mask = o_d < D + + if HAS_WEIGHT: + b_w = tl.load(w + i_g * D + o_d, mask=mask).to(tl.float32) + b_dw = tl.zeros((BD,), dtype=tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b_b = tl.load(b + i_g * D + o_d, mask=mask, other=0.0).to(tl.float32) + if HAS_BIAS: + b_db = tl.zeros((BD,), dtype=tl.float32) + + for i_t in range(i_sg * BS * G + i_g, min((i_sg * BS + BS) * G + i_g, T), G): + b_x = tl.load(x + i_t * D + o_d, mask=mask, other=0).to(tl.float32) + b_dy = tl.load(dy + i_t * D + o_d, mask=mask, other=0).to(tl.float32) + + if not IS_RMS_NORM: + b_mean = tl.load(mean + i_t) + b_rstd = tl.load(rstd + i_t) + # Compute dx + b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_xhat = tl.where(mask, b_xhat, 0.0) + if RECOMPUTE_OUTPUT: + b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat + if HAS_BIAS: + b_y = b_y + b_b + tl.store(y + i_t * D + o_d, b_y, mask=mask) + b_wdy = b_dy + if HAS_WEIGHT: + b_wdy = b_dy * b_w + b_dw += b_dy * b_xhat + if HAS_BIAS: + b_db += b_dy + if not IS_RMS_NORM: + b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D + b_c2 = tl.sum(b_wdy, axis=0) / D + b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd + else: + b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D + b_dx = (b_wdy - b_xhat * b_c1) * b_rstd + if HAS_DRESIDUAL: + b_dres = tl.load(dres + i_t * D + o_d, mask=mask, other=0).to(tl.float32) + b_dx += b_dres + # Write dx + b_dx = tl.cast(b_dx, dtype=dx.dtype.element_ty, fp_downcast_rounding='rtne') + if STORE_DRESIDUAL: + tl.store(dres_in + i_t * D + o_d, b_dx, mask=mask) + tl.store(dx + i_t * D + o_d, b_dx, mask=mask) + + if HAS_WEIGHT: + tl.store(dw + i_s * D + o_d, b_dw, mask=mask) + if HAS_BIAS: + tl.store(db + i_s * D + o_d, b_db, mask=mask) + + +def layer_norm_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, + num_groups: int = 1, +): + if residual is not None: + residual_dtype = residual.dtype + T, D, G = *x.shape, num_groups + if residual is not None: + assert residual.shape == (T, D) + if weight is not None: + assert weight.shape == (G * D,) + if bias is not None: + assert bias.shape == (G * D,) + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + res_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) + else: + res_out = None + mean = torch.empty((T,), dtype=torch.float, device=x.device) if not is_rms_norm else None + rstd = torch.empty((T,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + + if D <= 512: + NB = triton.cdiv(T, 2048) + def grid(meta): return (triton.cdiv(T, meta['BT']), ) + layer_norm_fwd_kernel[grid]( + x, + y, + weight, + bias, + residual, + res_out, + mean, + rstd, + eps, + T=T, + G=G, + D=D, + BD=BD, + NB=NB, + IS_RMS_NORM=is_rms_norm, + HAS_RESIDUAL=residual is not None, + STORE_RESIDUAL_OUT=res_out is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + else: + layer_norm_fwd_kernel1[(T,)]( + x, + y, + weight, + bias, + residual, + res_out, + mean, + rstd, + eps, + G=G, + D=D, + BD=BD, + IS_RMS_NORM=is_rms_norm, + HAS_RESIDUAL=residual is not None, + STORE_RESIDUAL_OUT=res_out is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + # res_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, res_out if res_out is not None else x + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + mean: torch.Tensor = None, + rstd: torch.Tensor = None, + dres: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, + num_groups: int = 1, +): + T, D, G = *x.shape, num_groups + assert dy.shape == (T, D) + if dres is not None: + assert dres.shape == (T, D) + if weight is not None: + assert weight.shape == (G * D,) + if bias is not None: + assert bias.shape == (G * D,) + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(T, D, dtype=x_dtype, device=x.device) + dres_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(T, D, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # each program handles one group only. + # cap per-group program count to T // G so no program is completely idle. + # without this, high-SM GPUs (e.g. B200, 160 SMs) with small T would + # launch idle programs whose make_block_ptr offsets exceed the tensor shape. + NS = min(triton.cdiv(get_multiprocessor_count(x.device.index), G), T // G) * G + BS = triton.cdiv(T, NS) + GS = NS // G + + dw = torch.empty((NS, D), dtype=torch.float, device=weight.device) if weight is not None else None + db = torch.empty((NS, D), dtype=torch.float, device=bias.device) if bias is not None else None + grid = (NS,) + + if D <= 512: + NB = triton.cdiv(T, 2048) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + y, + dy, + dx, + dw, + db, + dres, + dres_in, + mean, + rstd, + T=T, + G=G, + D=D, + BS=BS, + BD=BD, + NB=NB, + GS=GS, + IS_RMS_NORM=is_rms_norm, + HAS_DRESIDUAL=dres is not None, + STORE_DRESIDUAL=dres_in is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + else: + layer_norm_bwd_kernel1[grid]( + x, + weight, + bias, + y, + dy, + dx, + dw, + db, + dres, + dres_in, + mean, + rstd, + T=T, + G=G, + D=D, + BS=BS, + BD=BD, + GS=GS, + IS_RMS_NORM=is_rms_norm, + HAS_DRESIDUAL=dres is not None, + STORE_DRESIDUAL=dres_in is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + dw = dw.view(G, -1, D).sum(1).to(weight).view_as(weight) if weight is not None else None + db = db.view(G, -1, D).sum(1).to(bias).view_as(bias) if bias is not None else None + # Don't need to compute dres_in separately in this case + if has_residual and dx.dtype == x.dtype: + dres_in = dx + return (dx, dw, db, dres_in) if not recompute_output else (dx, dw, db, dres_in, y) + + +class LayerNormFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + weight, + bias, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1, + ): + x_shape_og = x.shape + + if x.shape[-1] % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + # reshape input data into 2D tensor + x = x.reshape(-1, (x.shape[-1] // num_groups)) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape_as(x) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + y, mean, rstd, res_out = layer_norm_fwd( + x, + weight, + bias, + eps, + residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + num_groups=num_groups, + ) + ctx.save_for_backward(res_out, weight, bias, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.num_groups = num_groups + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + y = y.reshape(x_shape_og) + return y if not prenorm else (y, res_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dy, *args): + x, weight, bias, mean, rstd = ctx.saved_tensors + dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, x.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dw, db, dresidual_in = layer_norm_bwd( + dy, + x, + weight, + bias, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + num_groups=ctx.num_groups, + ) + return ( + dx.reshape(ctx.x_shape_og), + dw, + db, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + None, + ) + + +def layer_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + ) + + +def group_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1, +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + num_groups, + ) + + +def rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + True, + ) + + +def layer_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1, +): + return LayerNormLinearFunction.apply( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + num_groups, + ) + + +def rms_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, +): + return layer_norm_linear( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True, + ) + + +def group_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1, +): + return layer_norm_linear( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=is_rms_norm, + num_groups=num_groups, + ) + + +class LayerNorm(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> LayerNorm: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class GroupNorm(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> GroupNorm: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return group_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=self.is_rms_norm, + num_groups=self.num_groups, + ) + + +class RMSNorm(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> RMSNorm: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return rms_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class LayerNormLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-5, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + num_groups=1, + ): + x_shape_og = x.shape + + if x.shape[-1] % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + # reshape input data into 2D tensor + x = x.reshape(-1, (x.shape[-1] // num_groups)) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape_as(x) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + y, mean, rstd, res_out = layer_norm_fwd( + x, + norm_weight, + norm_bias, + eps, + residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + num_groups=num_groups, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = linear_weight.to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(res_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.num_groups = num_groups + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, res_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, x.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd( + dy, + x, + norm_weight, + norm_bias, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + num_groups=ctx.num_groups, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y.view(-1, linear_weight.shape[-1])) + return ( + dx.reshape(ctx.x_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + None, + ) + + +class LayerNormLinear(nn.Module): + + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> LayerNormLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=False, + ) + + +class GroupNormLinear(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> GroupNormLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=self.is_rms_norm, + num_groups=self.num_groups, + ) + + +class RMSNormLinear(nn.Module): + + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> RMSNormLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True, + ) + + +class NormParallel(ParallelStyle): + + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False): + super().__init__() + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + def _replicate_module_fn( + self, name: str, module: nn.Module, device_mesh: DeviceMesh, + ): + for p_name, param in module.named_parameters(): + # simple replication with fixed ones_ init from LayerNorm/RMSNorm, which allow + # us to simply just use from_local + replicated_param = torch.nn.Parameter( + DTensor.from_local(param, device_mesh, [Replicate()], run_check=False), + ) + module.register_parameter(p_name, replicated_param) + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + input_tensor = inputs[0] + if isinstance(input_tensor, DTensor): + # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it + if input_tensor.placements != sequence_sharding: + input_tensor = input_tensor.redistribute( + placements=sequence_sharding, async_op=True, + ) + return input_tensor + elif isinstance(input_tensor, torch.Tensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + return DTensor.from_local( + input_tensor, device_mesh, sequence_sharding, run_check=False, + ) + else: + raise ValueError( + f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}", + ) + + @staticmethod + def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + self._replicate_module_fn, + partial(self._prepare_input_fn, self.sequence_sharding), + partial(self._prepare_output_fn, self.use_local_output), + ) diff --git a/fla/modules/layernorm_gated.py b/fla/modules/layernorm_gated.py new file mode 100644 index 0000000000000000000000000000000000000000..7702653c0832e49ff10871f3046fc522be7fe998 --- /dev/null +++ b/fla/modules/layernorm_gated.py @@ -0,0 +1,527 @@ +# Copyright (c) 2024, Tri Dao. +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + +from fla.utils import get_multiprocessor_count, input_guard + + +def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True): + dtype = x.dtype + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + z = z.float() if z is not None else z + if z is not None and not norm_before_gate: + x = x * F.silu(z) + if group_size is None: + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + else: + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) + out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight + if bias is not None: + out = out + bias + if z is not None and norm_before_gate: + out *= F.silu(z) + return out.to(dtype) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["B"] is not None, + "HAS_Z": lambda args: args["Z"] is not None, +}) +@triton.jit +def layer_norm_fwd_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + group = tl.program_id(1) + X += row * stride_x_row + group * N + Y += row * stride_y_row + group * N + if HAS_Z: + Z += row * stride_z_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + if HAS_BIAS: + B += group * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=cols < N).to(tl.float32) + x *= z * tl.sigmoid(z) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask).to(tl.float32) + y *= z * tl.sigmoid(z) + # Write output + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + z: torch.Tensor = None, + out: torch.Tensor = None, + group_size: int = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + assert x.stride(-1) == 1 + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + if out is not None: + assert out.shape == x.shape + else: + out = torch.empty_like(x) + assert out.stride(-1) == 1 + mean = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_N // 256, 1), 8) + grid = (M, ngroups) + layer_norm_fwd_kernel[grid]( + x, + out, + weight, + bias, + z, + mean, + rstd, + x.stride(0), + out.stride(0), + z.stride(0) if z is not None else 0, + M, + group_size, + eps, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps, + ) + return out, mean, rstd + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["B"] is not None, + "HAS_Z": lambda args: args["Z"] is not None, + "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None, +}) +@triton.jit +def layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DZ, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_z_row, + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dz_row, + stride_dw_row, + stride_db_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + rows_per_program, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + group = tl.program_id(1) + row_start = row_block_id * rows_per_program + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * stride_x_row + group * N + if HAS_Z: + Z += row_start * stride_z_row + group * N + DZ += row_start * stride_dz_row + group * N + DY += row_start * stride_dy_row + group * N + DX += row_start * stride_dx_row + group * N + if RECOMPUTE_OUTPUT: + Y += row_start * stride_y_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if (RECOMPUTE_OUTPUT or HAS_Z) and HAS_BIAS: + B += group * N + b = tl.load(B + cols, mask=mask, other=0.).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) + x_og = x + x = x_og * z * tl.sigmoid(z) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.) + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) + z_sigmoid = tl.sigmoid(z) + y = xhat * w + b if HAS_BIAS else xhat * w + if RECOMPUTE_OUTPUT: + tl.store(Y + cols, y * z * z_sigmoid, mask=mask) + dz = dy * y * z_sigmoid * (1 + z * (1 - z_sigmoid)) + tl.store(DZ + cols, dz, mask=mask) + dy *= z * z_sigmoid + else: + if RECOMPUTE_OUTPUT: + y = xhat * w + b if HAS_BIAS else xhat * w + tl.store(Y + cols, y, mask=mask) + wdy = w * dy + c1 = tl.sum(xhat * wdy, axis=0) / N + if not IS_RMS_NORM: + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + dx = (wdy - xhat * c1) * rstd + dw += dy * xhat + if HAS_BIAS: + db += dy + if HAS_Z and not NORM_BEFORE_GATE: + z_sigmoid = tl.sigmoid(z) + dz = dx * x_og * z_sigmoid * (1 + z * (1 - z_sigmoid)) + tl.store(DZ + cols, dz, mask=mask) + dx *= z * z_sigmoid + # Write dx + tl.store(DX + cols, dx, mask=mask) + + X += stride_x_row + if HAS_Z: + Z += stride_z_row + DZ += stride_dz_row + if RECOMPUTE_OUTPUT: + Y += stride_y_row + DY += stride_dy_row + DX += stride_dx_row + tl.store(DW + row_block_id * stride_dw_row + group * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * stride_db_row + group * N + cols, db, mask=mask) + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + z: torch.Tensor = None, + group_size: int = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, + recompute_output: bool = False, + dz: torch.Tensor = None, + out: torch.Tensor = None, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + assert x.stride(-1) == 1 + assert dy.stride(-1) == 1 + assert dy.shape == (M, N) + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + dx = torch.empty_like(x) + if dz is not None: + assert z is not None + assert dz.shape == z.shape + assert dz.stride(-1) == 1 + else: + dz = torch.empty_like(z) if z is not None else None + if recompute_output: + if out is None: + out = torch.empty_like(x) + assert out.shape == x.shape + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_N // 256, 1), 8) + sm_count = get_multiprocessor_count(x.device.index) + # If group size is small (e.g., 64), we're only using 1 warp. So having just 108 programs + # would limit the occupancy. + nrow_groups = math.ceil(sm_count * math.ceil(4 / num_warps) / ngroups) + _dw = torch.empty((nrow_groups, N), dtype=torch.float32, device=weight.device) + _db = torch.empty((nrow_groups, N), dtype=torch.float32, device=bias.device) if bias is not None else None + rows_per_program = math.ceil(M / nrow_groups) + grid = (nrow_groups, ngroups) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + z, + out if recompute_output else None, + dy, + dx, + _dw, + _db, + dz, + mean, + rstd, + x.stride(0), + z.stride(0) if z is not None else 0, + 0 if not recompute_output else out.stride(0), + dy.stride(0), + dx.stride(0), + dz.stride(0) if dz is not None else 0, + _dw.stride(0), + _db.stride(0) if _db is not None else 0, + M, group_size, eps, + rows_per_program, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps, + ) + dw = _dw.sum(0).to(weight.dtype) + db = _db.sum(0).to(bias.dtype) if bias is not None else None + return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out) + + +class LayerNormFn(torch.autograd.Function): + + @input_guard + @staticmethod + def forward(ctx, x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, + is_rms_norm=False): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + assert z.shape == x_shape_og + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y, mean, rstd = layer_norm_fwd( + x, + weight, + bias, + eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + ) + ctx.save_for_backward(x, weight, bias, mean, rstd, z) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.group_size = group_size + ctx.norm_before_gate = norm_before_gate + ctx.is_rms_norm = is_rms_norm + return y.reshape(x_shape_og) + + @input_guard + @staticmethod + def backward(ctx, dy): + x, weight, bias, mean, rstd, z = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + dx, dw, db, dz = layer_norm_bwd( + dy, + x, + weight, + bias, + ctx.eps, + mean, + rstd, + z, + ctx.group_size, + ctx.norm_before_gate, + ctx.is_rms_norm, + ) + dx = dx.reshape(ctx.x_shape_og) + dz = dz.reshape(ctx.x_shape_og) if dz is not None else None + return dx, dw, db, dz, None, None, None, None + + +def layernorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, is_rms_norm=False): + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm) + + +def rmsnorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True): + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, True) + + +class LayerNormGated(nn.Module): + + def __init__( + self, + hidden_size, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = True, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + torch.nn.init.zeros_(self.bias) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + return layernorm_fn(x, self.weight, self.bias, z=z, group_size=self.group_size, eps=self.eps, + norm_before_gate=self.norm_before_gate) + + +class RMSNormGated(nn.Module): + + def __init__( + self, + hidden_size, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = False, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + return rmsnorm_fn(x, self.weight, self.bias, z=z, eps=self.eps, group_size=self.group_size, + norm_before_gate=self.norm_before_gate) diff --git a/fla/modules/mlp.py b/fla/modules/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..74e9ab20b65c1d99c90ef82be60cf5ab433cd61b --- /dev/null +++ b/fla/modules/mlp.py @@ -0,0 +1,131 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from functools import partial +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from torch.distributed import DeviceMesh +from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.modules.activations import swiglu, swiglu_linear + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +class GatedMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + hidden_ratio: int | None = None, + intermediate_size: int | None = None, + hidden_act: str = 'swish', + fuse_swiglu: bool = True, + ) -> GatedMLP: + super().__init__() + + self.hidden_size = hidden_size + # the final number of params is `hidden_ratio * hidden_size^2` + # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio` + if hidden_ratio is None: + hidden_ratio = 4 + if intermediate_size is None: + intermediate_size = int(hidden_size * hidden_ratio * 2 / 3) + intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256) + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.fuse_swiglu = fuse_swiglu + + if hidden_act != 'swish': + raise ValueError(f'Unsupported hidden_act: {hidden_act}') + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + if self.fuse_swiglu: + self.swiglu_linear = SwiGLULinear() + + def forward( + self, + x: torch.Tensor, + **kwargs: Unpack[Any], + ) -> torch.Tensor: + gate, y = self.gate_proj(x), self.up_proj(x) + if self.fuse_swiglu: + return self.swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias) + else: + return self.down_proj(swiglu(gate, y)) + + +class SwiGLULinear(nn.Module): + + def forward(self, x, y, weight, bias): + return swiglu_linear(x, y, weight, bias) + + +class SwiGLULinearParallel(ParallelStyle): + def __init__( + self, + *, + input_layouts: Placement | None = None, + output_layouts: Placement | None = None, + use_local_output: bool = True, + ): + super().__init__() + self.input_layouts = (input_layouts or Shard(-1),) + self.output_layouts = (output_layouts or Replicate(),) + self.desired_input_layouts = (Shard(-1),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn( + input_layouts, desired_input_layouts, mod, inputs, device_mesh, + ): + x, y, weight, bias = inputs + if not isinstance(x, DTensor): + x = DTensor.from_local(x, device_mesh, input_layouts, run_check=False) + if x.placements != desired_input_layouts: + x = x.redistribute(placements=desired_input_layouts, async_op=True) + + if not isinstance(y, DTensor): + y = DTensor.from_local(y, device_mesh, input_layouts, run_check=False) + if y.placements != desired_input_layouts: + y = y.redistribute(placements=desired_input_layouts, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, (Shard(1),)) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, (Replicate(),)) + + return x, y, weight, bias + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + # Rowwise sharding produces partial output, depending on output layouts: + # 1. to replicate -> allreduce + # 2. to shard -> reduce_scatter + if outputs.placements != output_layouts: + outputs = outputs.redistribute(placements=output_layouts, async_op=True) + # back to local tensor if use_local_output is True + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts), + output_fn=partial(self._prepare_output_fn, self.output_layouts, self.use_local_output), + ) diff --git a/fla/modules/parallel.py b/fla/modules/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..f889326167ffdec4aef6178a7ff6344fb442d8ee --- /dev/null +++ b/fla/modules/parallel.py @@ -0,0 +1,40 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch.nn as nn +from torch.distributed import DeviceMesh +from torch.distributed.tensor import distribute_module +from torch.distributed.tensor.parallel import ParallelStyle +from torch.distributed.tensor.placement_types import Placement + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + + +class PrepareModuleWeight(ParallelStyle): + def __init__(self, *, layouts: Placement | None = None): + super().__init__() + self.layouts = layouts + + def _replicate_module_fn( + self, + name: str, + module: nn.Module, + device_mesh: DeviceMesh, + ): + for p_name, param in module.named_parameters(): + replicated_param = nn.Parameter( + DTensor.from_local(param, device_mesh, [self.layouts], run_check=False), + ) + module.register_parameter(p_name, replicated_param) + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=self._replicate_module_fn, + input_fn=None, + output_fn=None, + ) diff --git a/fla/modules/rotary.py b/fla/modules/rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..6f43be7bd6eb3a2e288421a7c9da37826b5c5a98 --- /dev/null +++ b/fla/modules/rotary.py @@ -0,0 +1,511 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from einops import rearrange, repeat + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2) + + +def rotary_embedding_ref(x, cos, sin, interleaved=False): + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat(cos, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + sin = repeat(sin, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], -1) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['B', 'H', 'D', 'INTERLEAVED'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def rotary_embedding_kernel( + x, + cos, + sin, + y, + cu_seqlens, + chunk_indices, + seq_offsets, + T, + B: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + TR: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, + IS_VARLEN: tl.constexpr, + INTERLEAVED: tl.constexpr, + CONJUGATE: tl.constexpr, +): + i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1) + T = eos - bos + x = x + bos * H*D + i_h * D + y = y + bos * H*D + i_h * D + else: + i_n = i_b + x = x + i_n * T*H*D + i_h * D + y = y + i_n * T*H*D + i_h * D + + if i_t * BT >= T: + return + + o_t = i_t * BT + tl.arange(0, BT) + if not IS_SEQLEN_OFFSETS_TENSOR: + o_cs = o_t + seq_offsets + else: + o_cs = o_t + tl.load(seq_offsets + i_n) + m_t = (o_t >= 0) & (o_t < T) & (o_cs >= 0) & (o_cs < TR) + + if not INTERLEAVED: + # Load the 1st and 2nd halves of x, do calculation, then store to 1st and 2nd halves of out + o_r = tl.arange(0, BD // 2) + p_x = x + o_t[:, None] * H*D + o_r[None, :] + p_cos = cos + (o_cs[:, None] * R + o_r[None, :]) + p_sin = sin + (o_cs[:, None] * R + o_r[None, :]) + mask = m_t[:, None] & (o_r < R)[None, :] + + b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32) + b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32) + b_x0 = tl.load(p_x, mask=mask, other=0.0).to(tl.float32) + b_x1 = tl.load(p_x + R, mask=mask, other=0.0).to(tl.float32) + if CONJUGATE: + b_sin = -b_sin + b_o0 = b_x0 * b_cos - b_x1 * b_sin + b_o1 = b_x0 * b_sin + b_x1 * b_cos + # write back result + p_y = y + (o_t[:, None] * H*D + o_r[None, :]) + tl.store(p_y, b_o0, mask=mask) + tl.store(p_y + R, b_o1, mask=mask) + else: + # We don't want to load x[0, 2, 4, ...] and x[1, 3, 5, ...] separately since both are slow. + # Instead, we load x0 = x[0, 1, 2, 3, ...] and x1 = x[1, 0, 3, 2, ...]. + # Loading x0 will be fast but x1 will be slow. + # Then we load cos = cos[0, 0, 1, 1, ...] and sin = sin[0, 0, 1, 1, ...]. + # Then we do the calculation and use tl.where to pick put the right outputs for the even + # and for the odd indices. + o_d = tl.arange(0, BD) + o_d_swap = o_d + ((o_d + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ... + o_d_repeat = tl.arange(0, BD) // 2 + p_x0 = x + o_t[:, None] * H*D + o_d[None, :] + p_x1 = x + o_t[:, None] * H*D + o_d_swap[None, :] + p_cos = cos + (o_cs[:, None] * R + o_d_repeat[None, :]) + p_sin = sin + (o_cs[:, None] * R + o_d_repeat[None, :]) + mask = m_t[:, None] & (o_d_repeat < R)[None, :] + + b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32) + b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32) + b_x0 = tl.load(p_x0, mask=mask, other=0.0).to(tl.float32) + b_x1 = tl.load(p_x1, mask=mask, other=0.0).to(tl.float32) + if CONJUGATE: + b_sin = -b_sin + b_o0 = b_x0 * b_cos + b_o1 = b_x1 * b_sin + b_y = tl.where(o_d[None, :] % 2 == 0, b_o0 - b_o1, b_o0 + b_o1) + p_y = y + (o_t[:, None] * H*D + o_d[None, :]) + tl.store(p_y, b_y, mask=mask) + + +def rotary_embedding_fwdbwd( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + seqlen_offsets: int | torch.Tensor = 0, + cu_seqlens: torch.Tensor | None = None, + interleaved: bool = False, + inplace: bool = False, + conjugate: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + """ + Args: + x: [B, T, H, D]. + cos: [TR, R / 2] + sin: [TR, R / 2] + seqlen_offsets: integer or integer tensor of size [N] + cu_seqlens: [N + 1,] or None + + Returns: + y: [B, T, H, D] + """ + is_varlen = cu_seqlens is not None + + B, T, H, D = x.shape + N = B if not is_varlen else cu_seqlens.shape[0] - 1 + TR, R = cos.shape + R2 = R * 2 + + assert D <= 256, "Only support D <= 256" + assert TR >= T, f"TR must be >= T, got {TR} and {T}" + + assert cos.dtype == sin.dtype, f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}" + assert x.dtype == cos.dtype, f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}" + + if isinstance(seqlen_offsets, torch.Tensor): + assert seqlen_offsets.shape == (N,) + assert seqlen_offsets.dtype in [torch.int32, torch.int64] + else: + assert seqlen_offsets + T <= TR + + y = torch.empty_like(x) if not inplace else x + if R2 < D and not inplace: + y[..., R2:].copy_(x[..., R2:]) + + BD = triton.next_power_of_2(R2) + BT = min(128, triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index)))) + if chunk_indices is None and is_varlen: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) if is_varlen else triton.cdiv(T, BT) + + grid = (NT, B, H) + rotary_embedding_kernel[grid]( + x, + cos, + sin, + y, + cu_seqlens, + chunk_indices, + seqlen_offsets, + B=B, + T=T, + H=H, + D=D, + R=R, + TR=TR, + BT=BT, + BD=BD, + IS_SEQLEN_OFFSETS_TENSOR=isinstance(seqlen_offsets, torch.Tensor), + IS_VARLEN=is_varlen, + INTERLEAVED=interleaved, + CONJUGATE=conjugate, + ) + return y + + +class RotaryEmbeddingFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + cos, + sin, + interleaved=False, + inplace=False, + seqlen_offsets: int | torch.Tensor = 0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.LongTensor | None = None, + ): + y = rotary_embedding_fwdbwd( + x, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + interleaved=interleaved, + inplace=inplace, + chunk_indices=chunk_indices, + ) + if isinstance(seqlen_offsets, int): + # Can't save int with save_for_backward + ctx.save_for_backward(cos, sin, cu_seqlens) + ctx.seqlen_offsets = seqlen_offsets + else: + ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) + ctx.seqlen_offsets = None + ctx.interleaved = interleaved + ctx.inplace = inplace + ctx.chunk_indices = chunk_indices + return y if not inplace else x + + @staticmethod + @input_guard + def backward(ctx, do): + seqlen_offsets = ctx.seqlen_offsets + if seqlen_offsets is None: + cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors + else: + cos, sin, cu_seqlens = ctx.saved_tensors + # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with + # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works. + if not ctx.interleaved and not ctx.inplace: + do = do.clone() + dx = rotary_embedding_fwdbwd( + do, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + interleaved=ctx.interleaved, + inplace=ctx.inplace, + conjugate=True, + chunk_indices=ctx.chunk_indices, + ) + return dx, None, None, None, None, None, None, None + + +def rotary_embedding( + x, + cos, + sin, + interleaved=False, + inplace=False, + seqlen_offsets: int | torch.Tensor = 0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + """ + Args: + x: [B, T, H, D] + cos, sin: [TR, R//2] + interleaved: + If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). + inplace: + If True, apply rotary embedding in-place. + seqlen_offsets: [N,] or int. + Each sequence in x is shifted by this amount. + Most commonly used in inference when we have KV cache. + cu_seqlens: [N + 1,] or None + + Returns: + out: [B, T, H, D] + """ + return RotaryEmbeddingFunction.apply( + x, + cos, + sin, + interleaved, + inplace, + seqlen_offsets, + cu_seqlens, + chunk_indices, + ) + + +class RotaryEmbedding(nn.Module): + """ + The rotary position embeddings from RoFormer_ (Su et. al). + A crucial insight from the method is that the query and keys are + transformed by rotation matrices which depend on the relative positions. + + Other implementations are available in the Rotary Transformer repo_ and in + GPT-NeoX_, GPT-NeoX was an inspiration + + .. _RoFormer: https://arxiv.org/abs/2104.09864 + .. _repo: https://github.com/ZhuiyiTechnology/roformer + .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox + + If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). + A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 + Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py + """ + + def __init__( + self, + dim: int, + base: float = 10000.0, + scale_base: float | None = None, + interleaved: bool = False, + pos_idx_in_fp32: bool = True, + device: torch.device | None = None, + ): + """ + interleaved: + If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). + pos_idx_in_fp32: + If True, the position indices [0.0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. + This option was added because previously (before 2023-07-02), when we construct + the position indices, we use the dtype of self.inv_freq. + In most cases this would be fp32, but if the model is trained in pure bf16 (not mixed precision), then + self.inv_freq would be bf16, and the position indices are also in bf16. + Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the + embeddings for some positions will coincide. + To maintain compatibility with models previously trained in pure bf16, we add this option. + """ + super().__init__() + + self.dim = dim + self.base = float(base) + self.scale_base = scale_base + self.interleaved = interleaved + self.pos_idx_in_fp32 = pos_idx_in_fp32 + self.device = device + + # Generate and save the inverse frequency buffer (non trainable) + self.register_buffer("inv_freq", torch.empty(-(dim // -2), dtype=torch.float32, device=device), persistent=False) + + scale = None + if scale_base is not None: + scale = torch.empty(-(dim // -2), dtype=torch.float32, device=device) + self.register_buffer("scale", scale, persistent=False) + + self._seq_len_cached = 0 + self._cos_cached = None + self._sin_cached = None + self._cos_k_cached = None + self._sin_k_cached = None + + self.reset_parameters() + + def reset_parameters(self): + with torch.no_grad(): + self.inv_freq.copy_(self._compute_inv_freq(device=self.inv_freq.device)) + if self.scale_base is not None: + self.scale.copy_(self._compute_scale(device=self.scale.device)) + + def __repr__(self): + s = f"{self.__class__.__name__}(" + s += f"dim={self.dim}, " + s += f"base={self.base}, " + s += f"interleaved={self.interleaved}, " + if self.scale_base is not None: + s += f"scale_base={self.scale_base}, " + s += f"pos_idx_in_fp32={self.pos_idx_in_fp32})" + return s + + def _compute_inv_freq(self, device=None): + return 1.0 / ( + self.base + ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) + ) + + def _compute_scale(self, device=None): + return (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + 0.4 * self.dim) / (1.4 * self.dim) + + def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): + # Reset the tables if the sequence length has changed, + # if we're on a new device (possibly due to tracing for instance), + # or if we're switching from inference mode to training + if ( + seqlen > self._seq_len_cached + or self._cos_cached is None + or self._cos_cached.device != device + or self._cos_cached.dtype != dtype + or (self.training and self._cos_cached.is_inference()) + ): + self._seq_len_cached = seqlen + # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 + # And the output of arange can be quite large, so bf16 would lose a lot of precision. + # However, for compatibility reason, we add an option to use the dtype of self.inv_freq. + if self.pos_idx_in_fp32: + t = torch.arange(seqlen, device=device, dtype=torch.float32) + # We want fp32 here as well since inv_freq will be multiplied with t, and the output + # will be large. Having it in bf16 will lose a lot of precision and cause the + # cos & sin output to change significantly. + # We want to recompute self.inv_freq if it was not loaded in fp32 + if self.inv_freq.dtype != torch.float32: + inv_freq = self._compute_inv_freq(device=device) + else: + inv_freq = self.inv_freq + else: + t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) + inv_freq = self.inv_freq + # Don't do einsum, it converts fp32 to fp16 under AMP + # freqs = torch.einsum("i,j->ij", t, self.inv_freq) + freqs = torch.outer(t, inv_freq) + if self.scale is None: + self._cos_cached = torch.cos(freqs).to(dtype) + self._sin_cached = torch.sin(freqs).to(dtype) + else: + power = ( + torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) + - seqlen // 2 + ) / self.scale_base + scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") + # We want the multiplication by scale to happen in fp32 + self._cos_cached = (torch.cos(freqs) * scale).to(dtype) + self._sin_cached = (torch.sin(freqs) * scale).to(dtype) + self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) + self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + seqlen_offset: int | torch.Tensor = 0, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, + chunk_indices: torch.LongTensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + q: [B, T, H, D] + k: [B, T, H, D] + seqlen_offset: + [N] or int. + Each sequence in x is shifted by this amount. + Most commonly used in inference when we have KV cache. + cu_seqlens: [N + 1] or None + max_seqlen: int + """ + if max_seqlen is not None: + self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype) + elif isinstance(seqlen_offset, int): + self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype) + if self.scale is None: + q = rotary_embedding( + q, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + k = rotary_embedding( + k, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + else: + q = rotary_embedding( + q, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + k = rotary_embedding( + k, + self._cos_k_cached, + self._sin_k_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + return q, k diff --git a/fla/modules/token_shift.py b/fla/modules/token_shift.py new file mode 100644 index 0000000000000000000000000000000000000000..8748d32890374d4a32105c61aeb3530fabb28473 --- /dev/null +++ b/fla/modules/token_shift.py @@ -0,0 +1,552 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard, tensor_cache + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +def token_shift_ref( + x: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, +) -> torch.Tensor: + if cu_seqlens is not None: + # Variable length mode with cu_seqlens + assert x.dim() == 3, "Input must be [B, T, D]" + B, T, D = x.shape + assert B == 1, "Batch size must be 1 when using cu_seqlens" + + result = torch.zeros_like(x) + N = cu_seqlens.shape[0] - 1 + + for i in range(N): + start = cu_seqlens[i].item() + end = cu_seqlens[i+1].item() + seq_len = end - start + + if seq_len <= 1: + # For sequences of length 1 or 0, delta is simply -x + result[0, start:end] = -x[0, start:end] + else: + # For longer sequences, handle padding manually + shifted = torch.zeros_like(x[0, start:end]) + shifted[1:] = x[0, start:end-1] + delta = shifted - x[0, start:end] + result[0, start:end] = delta + + return result + else: + time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1)) + shifted = time_shift(x) + delta = shifted - x + return delta + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + ], + key=['BD'], + **autotune_cache_kwargs, +) +@triton.jit +def token_shift_fwd_kernel_short( + x, + y, + cu_seqlens, + cache, + cache_out, + T, + D: tl.constexpr, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_DECODE: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = i_b + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + g_t = i_t + bos + + if g_t >= eos: + return + + is_first_pos = (i_t == 0) + is_last_pos = (g_t == eos - 1) + else: + g_t = i_t + is_first_pos = (g_t == 0) + is_last_pos = (g_t == T - 1) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + if IS_VARLEN: + base_offset = g_t * D + o_d + else: + base_offset = i_b * T*D + g_t * D + o_d + + b_x = tl.load(x + base_offset, mask=m_d) + if IS_VARLEN: + cache_offset = i_n * D + o_d # i_n is seq index + else: + cache_offset = i_b * D + o_d # i_b is batch index + + if IS_DECODE and USE_INITIAL_STATE: + b_cache = tl.load(cache + cache_offset, mask=m_d) + delta = b_cache - b_x + tl.store(y + base_offset, delta, mask=m_d) + if STORE_FINAL_STATE: + tl.store(cache_out + cache_offset, b_x, mask=m_d) + return + + if is_first_pos: + # First position in sequence: delta = -hidden_states + if USE_INITIAL_STATE: + # cache shape: [N, D] + b_cache = tl.load(cache + cache_offset, mask=m_d) + delta = b_cache - b_x + tl.store(y + base_offset, delta, mask=m_d) + else: + tl.store(y + base_offset, -b_x, mask=m_d) + return + + # Other positions: delta = prev - curr + if IS_VARLEN: + prev_offset = (g_t-1) * D + o_d + else: + prev_offset = i_b * T*D + (g_t-1) * D + o_d + + prev_values = tl.load(x + prev_offset, mask=m_d) + delta = prev_values - b_x + tl.store(y + base_offset, delta, mask=m_d) + if STORE_FINAL_STATE: + if is_last_pos: + tl.store(cache_out + cache_offset, b_x, mask=m_d) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + ], + key=['BD', 'NB'], + **autotune_cache_kwargs, +) +@triton.jit +def token_shift_fwd_kernel_long( + x, + y, + cu_seqlens, + chunk_indices, + cache, + cache_out, + T, + D: tl.constexpr, + BD: tl.constexpr, + BT: tl.constexpr, + NB: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, +): + i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), \ + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1) + t_start = i_t * BT + t_end = tl.minimum(t_start + BT, eos - bos) + else: + i_n = i_b + bos, eos = i_b * T, (i_b + 1) * T + t_start = i_t * BT + t_end = tl.minimum(t_start + BT, T) + + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + + for t in range(t_start, t_end): + global_t = bos + t + offset = global_t * D + o_d + b_x = tl.load(x + offset, mask=m_d) + is_first = (global_t == bos) + if is_first: + if USE_INITIAL_STATE: + # cache shape: [N, D] + cache_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d + b_cache = tl.load(cache + cache_off, mask=m_d) + delta = b_cache - b_x + else: + delta = -b_x + else: + prev_off = offset - D + b_prev = tl.load(x + prev_off, mask=m_d) + delta = b_prev - b_x + + tl.store(y + offset, delta, mask=m_d) + + if STORE_FINAL_STATE: + if global_t == eos - 1: + cache_out_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d + tl.store(cache_out + cache_out_off, b_x, mask=m_d) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['grad_cache_out'] is not None, + 'HAS_DCACHE': lambda args: args['grad_cache_in'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + ], + key=['BD'], + **autotune_cache_kwargs, +) +@triton.jit +def token_shift_bwd_kernel_short( + dx, + dy, + cu_seqlens, + grad_cache_in, + grad_cache_out, + T, + D: tl.constexpr, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + HAS_DCACHE: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = i_b + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + g_t = i_t + bos + if g_t >= eos: + return + is_first_pos = (g_t == bos) + is_last_pos = (g_t == eos - 1) + else: + g_t = i_t + is_first_pos = (g_t == 0) + is_last_pos = (g_t == T - 1) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + if IS_VARLEN: + base_offset = g_t * D + o_d + # This should not be used for varlen + cache_off = i_n * D + o_d + else: + base_offset = i_b * T * D + g_t * D + o_d + cache_off = i_b * D + o_d + + b_dy = tl.load(dy + base_offset, mask=m_d) + + if is_last_pos: + # grad = -grad_delta[t] + grad_cache_in（from next rank） + if HAS_DCACHE: + b_dy_cache = tl.load(grad_cache_in + cache_off, mask=m_d) + b_dx = -b_dy + b_dy_cache + else: + b_dx = -b_dy + else: + # grad = -grad_delta[t] + grad_delta[t+1] + if IS_VARLEN: + next_offset = (g_t + 1) * D + o_d + else: + next_offset = i_b * T * D + (g_t + 1) * D + o_d + b_dx = -b_dy + tl.load(dy + next_offset, mask=m_d) + + tl.store(dx + base_offset, b_dx, mask=m_d) + + if USE_INITIAL_STATE: + if is_first_pos: + tl.store(grad_cache_out + cache_off, b_dy, mask=m_d) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['grad_cache_out'] is not None, + 'HAS_DCACHE': lambda args: args['grad_cache_in'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + ], + key=['BD', 'NB'], + **autotune_cache_kwargs, +) +@triton.jit +def token_shift_bwd_kernel_long( + dx, + dy, + cu_seqlens, + chunk_indices, + grad_cache_in, + grad_cache_out, + T, + D: tl.constexpr, + BD: tl.constexpr, + BT: tl.constexpr, + NB: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + HAS_DCACHE: tl.constexpr, +): + i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ + tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1) + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, eos - bos) + else: + bos, eos = i_b * T, (i_b + 1) * T + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, T) + + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + cache_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d + + for t in range(t_start, t_end): + global_t = bos + t + offset = global_t * D + o_d + b_dy = tl.load(dy + offset, mask=m_d) + + if global_t == eos - 1: + if HAS_DCACHE: + b_dy_cache = tl.load(grad_cache_in + cache_off, mask=m_d) + b_dx = -b_dy + b_dy_cache + else: + b_dx = -b_dy + else: + next_off = offset + D + b_dx = -b_dy + tl.load(dy + next_off, mask=m_d) + + tl.store(dx + offset, b_dx, mask=m_d) + + if USE_INITIAL_STATE: + if global_t == bos: + tl.store(grad_cache_out + cache_off, b_dy, mask=m_d) + + +@tensor_cache +def prepare_maxlens(cu_seqlens: torch.LongTensor) -> int: + return torch.max(cu_seqlens.diff()).item() + + +def token_shift_fwd( + x: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + cache: torch.Tensor | None = None, + output_cache: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, D = x.shape + y = torch.empty_like(x) + use_short_kernel = T <= 4096 + + if cu_seqlens is not None: + T = prepare_maxlens(cu_seqlens) + N = len(cu_seqlens) - 1 + else: + N = B + + if output_cache: + cache_out = torch.empty((N, D), device=x.device, dtype=x.dtype) + else: + cache_out = None + + if use_short_kernel: + if cu_seqlens is not None: + N = len(cu_seqlens) - 1 + else: + N = B + BD = triton.next_power_of_2(D) + grid = (N, T) + IS_DECODE = T == 1 or (B == 1 and T == N) + token_shift_fwd_kernel_short[grid]( + x=x, + y=y, + cu_seqlens=cu_seqlens, + cache=cache, + cache_out=cache_out, + T=T, + D=D, + BD=BD, + STORE_FINAL_STATE=output_cache, + IS_DECODE=IS_DECODE, + ) + else: + BT = min(64, triton.next_power_of_2(triton.cdiv(max(16, B*T), get_multiprocessor_count(x.device.index)))) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + + BD = triton.next_power_of_2(D) + NB = triton.cdiv(B*T, 1024) + + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, N) + token_shift_fwd_kernel_long[grid]( + x, + y, + cu_seqlens, + chunk_indices, + cache, + cache_out, + T, + D=D, + BD=BD, + BT=BT, + NB=NB, + STORE_FINAL_STATE=output_cache, + ) + + return y, N, T, use_short_kernel, cache_out + + +def token_shift_bwd( + dy: torch.Tensor, + N: int, + T: int, + dcache: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + use_short_kernel: bool = True, + has_init_cache: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + D = dy.shape[2] + BD = triton.next_power_of_2(D) + dx = torch.empty_like(dy) + if has_init_cache: + grad_cache_out = torch.empty((N, D), device=dy.device, dtype=dy.dtype) + else: + grad_cache_out = None + if use_short_kernel: + grid = (N, T) + token_shift_bwd_kernel_short[grid]( + dy=dy, + dx=dx, + cu_seqlens=cu_seqlens, + grad_cache_in=dcache, + grad_cache_out=grad_cache_out, + T=T, + D=D, + BD=BD, + ) + else: + BT = min(64, triton.next_power_of_2(triton.cdiv(max(16, dy.numel() // D), + get_multiprocessor_count(dy.device.index)))) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + NB = triton.cdiv(N * dy.shape[1], 1024) + BD = triton.next_power_of_2(D) + + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, N) + token_shift_bwd_kernel_long[grid]( + dx, + dy, + cu_seqlens, + chunk_indices, + dcache, + grad_cache_out, + T, + D=D, + BD=BD, + BT=BT, + NB=NB, + ) + return dx, grad_cache_out + + +class TokenShift(torch.autograd.Function): + + @staticmethod + @input_guard + def forward(ctx, x: torch.Tensor, cu_seqlens: torch.Tensor | None = None, + cache: torch.Tensor | None = None, output_cache: bool = False, + chunk_indices: torch.LongTensor | None = None): + output, N, T, use_short_kernel, cache_out = token_shift_fwd(x, cu_seqlens, cache, output_cache, chunk_indices) + ctx.cu_seqlens = cu_seqlens + ctx.chunk_indices = chunk_indices + ctx.N = N + ctx.T = T + ctx.use_short_kernel = use_short_kernel + ctx.has_cache = cache is not None + return output, cache_out + + @staticmethod + @input_guard + def backward(ctx, dy: torch.Tensor, dcache: torch.Tensor | None = None): + dx, grad_cache = token_shift_bwd(dy, ctx.N, ctx.T, dcache, ctx.cu_seqlens, + ctx.use_short_kernel, ctx.has_cache, ctx.chunk_indices) + return dx, None, grad_cache, None, None + + +def token_shift( + x: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + cache: torch.Tensor | None = None, + output_cache: bool = False, + chunk_indices: torch.LongTensor | None = None, +): + """ + Token-shift operation implemented with Triton kernels. + + Args: + x: Input tensor of shape [B, T, D] (or [1, T, D] when `cu_seqlens` is supplied). + cu_seqlens: Optional cumulative sequence lengths of shape [B + 1]. + When supplied, `x.shape[0]` must be 1 and `x.dim()` must be 3. + cache: Optional cache tensor of shape [N, D] that holds the last token + from the previous call. + output_cache: Whether to return the updated cache alongside the output. + In previous versions this parameter did not exist and the + cache was always dropped; to preserve backward compatibility + the default is False. + + Returns: + output: Tensor of shape [B, T, D] after applying the token-shift. + + cache_out: Tensor of shape [B, 1, D] containing the last token that + should be fed as `cache` in the next call. Only returned + when `output_cache=True`. + """ + if cu_seqlens is not None: + assert x.dim() == 3, "Input must be [B, T, D]" + assert x.shape[0] == 1, "Batch size must be 1 when using cu_seqlens" + + output, cache_out = TokenShift.apply(x, cu_seqlens, cache, output_cache, chunk_indices) + if output_cache: + return output, cache_out + else: + return output diff --git a/fla/ops/__init__.py b/fla/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..32afc2f0fa9480606a327157187ba7def86f7915 --- /dev/null +++ b/fla/ops/__init__.py @@ -0,0 +1,78 @@ + +from .abc import chunk_abc +from .attn import parallel_attn +from .based import fused_chunk_based, parallel_based +from .comba import chunk_comba, fused_recurrent_comba +from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule +from .forgetting_attn import parallel_forgetting_attn +from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn +from .generalized_delta_rule import ( + chunk_dplr_delta_rule, + chunk_iplr_delta_rule, + fused_recurrent_dplr_delta_rule, + fused_recurrent_iplr_delta_rule, +) +from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla +from .gsa import chunk_gsa, fused_recurrent_gsa +from .hgrn import fused_recurrent_hgrn +from .kda import chunk_kda, fused_recurrent_kda +from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn +from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn +from .log_linear_attn import chunk_log_linear_attn +from .mesa_net import chunk_mesa_net +from .nsa import parallel_nsa +from .path_attn import parallel_path_attn +from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention +from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 +from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 +from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla + +__all__ = [ + 'chunk_abc', + 'chunk_comba', + 'chunk_delta_rule', + 'chunk_dplr_delta_rule', + 'chunk_gated_delta_rule', + 'chunk_gdn', + 'chunk_gla', + 'chunk_gsa', + 'chunk_iplr_delta_rule', + 'chunk_kda', + 'chunk_lightning_attn', + 'chunk_linear_attn', + 'chunk_log_linear_attn', + 'chunk_mesa_net', + 'chunk_retention', + 'chunk_rwkv6', + 'chunk_rwkv7', + 'chunk_simple_gla', + 'fused_chunk_based', + 'fused_chunk_delta_rule', + 'fused_chunk_gla', + 'fused_chunk_linear_attn', + 'fused_chunk_retention', + 'fused_chunk_simple_gla', + 'fused_recurrent_comba', + 'fused_recurrent_delta_rule', + 'fused_recurrent_dplr_delta_rule', + 'fused_recurrent_gated_delta_rule', + 'fused_recurrent_gdn', + 'fused_recurrent_gla', + 'fused_recurrent_gsa', + 'fused_recurrent_hgrn', + 'fused_recurrent_iplr_delta_rule', + 'fused_recurrent_kda', + 'fused_recurrent_lightning_attn', + 'fused_recurrent_linear_attn', + 'fused_recurrent_retention', + 'fused_recurrent_rwkv6', + 'fused_recurrent_rwkv7', + 'fused_recurrent_simple_gla', + 'parallel_attn', + 'parallel_based', + 'parallel_forgetting_attn', + 'parallel_nsa', + 'parallel_path_attn', + 'parallel_retention', + 'parallel_simple_gla', +] diff --git a/fla/ops/abc/__init__.py b/fla/ops/abc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..48d5c6d8e9a1ba557195486ef42aefe5e8fd79df --- /dev/null +++ b/fla/ops/abc/__init__.py @@ -0,0 +1,6 @@ + +from .chunk import chunk_abc + +__all__ = [ + 'chunk_abc', +] diff --git a/fla/ops/abc/chunk.py b/fla/ops/abc/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..a07228ef8da30fb5b982aa4526c091447d055895 --- /dev/null +++ b/fla/ops/abc/chunk.py @@ -0,0 +1,1115 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import softmax_bwd, softmax_fwd +from fla.ops.utils.logcumsumexp import logcumsumexp_fwd_kernel +from fla.ops.utils.op import exp +from fla.utils import input_guard + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_fwd_kernel_h( + k, + v, + z, + h, + h0, + ht, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, + NORMK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, +): + i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h = tl.make_block_ptr(h0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + if NORMK: + p_z0 = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_k * BK,), (BK,), (0,)) + else: + p_z0 = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_v * BV,), (BV,), (0,)) + b_zp = tl.load(p_z0).to(tl.float32) + for i_t in range(NT): + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + if NORMK: + p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + b_r, b_zp = exp(b_zp - b_zc), b_zc + # [BK, BV] + b_h = b_h * b_r[:, None] + b_k = exp(b_k - b_zc[:, None]).to(b_k.dtype) + else: + p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) + # [BV,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + b_r, b_zp = exp(b_zp - b_zc), b_zc + # [BK, BV] + b_h = b_h * b_r[None, :] + b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) + # [BK, BV] + b_h += tl.dot(b_k, b_v, allow_tf32=False) + + if STORE_FINAL_STATE: + p_h = tl.make_block_ptr(ht + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_fwd_kernel_intra_K( + v, + z, + o, + A, + T, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i = i_c // NC, i_c % NC + + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) + # [BV,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BV] + b_o = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(0, i_i): + p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_o += tl.dot(b_A, exp(b_v - b_zn[None, :]).to(b_v.dtype), allow_tf32=False) + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_o *= exp(b_zn[None, :] - b_z) + + o_i = tl.arange(0, BC) + o_A = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, BC): + p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) + # [BC,] + b_A = tl.load(A + o_A + j, mask=m_A, other=0) + # [BV,] + b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) + # [BC, BV] + # avoid 0 * inf = inf + m_i = o_i[:, None] >= j + b_o += tl.where(m_i, b_A[:, None] * exp(b_v[None, :] - b_z), 0) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_fwd_kernel_K( + q, + k, + z, + h, + o, + A, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_p = tl.maximum(i_t * BT - 1, 0) + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_o += tl.dot(b_q, b_h, allow_tf32=False) + # [BT, BT] + b_A += tl.dot(b_q, b_k, allow_tf32=False) + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_z = tl.load(p_z, boundary_check=(0, 1)) + # [BT, BV] + p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) + b_zp = tl.load(p_zp, boundary_check=(0,)) + b_o = b_o * exp(b_zp[None, :] - b_z) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BT] + b_A = tl.where(m_s, b_A, 0.) + if i_v == 0: + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_fwd_kernel_intra_V( + q, + k, + z, + A, + scale, + T, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, +): + i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + n_bh = tl.num_programs(2) + + if i_i > i_j: + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_q = (b_q * exp(b_zn[None, :] - b_z) * scale).to(b_q.dtype) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = exp(b_k - b_zn[:, None]).to(b_k.dtype) + # [BC, BC] + b_A = tl.dot(b_q, b_k, allow_tf32=False) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + elif i_i == i_j: + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_j * BC) * K + i_k * BK,), (BK,), (0,)) + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_z = tl.load(p_z, boundary_check=(0, 1)) + + o_i = tl.arange(0, BC) + o_A = (i_bh + i_k * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, BC): + # [BK,] + b_k = tl.load(p_k, boundary_check=(0,)).to(tl.float32) + # [BC,] + b_A = tl.sum(b_q * exp(b_k[None, :] - b_z) * scale, 1) + b_A = tl.where(o_i >= j, b_A, 0.) + tl.store(A + o_A + j, b_A.to(b_q.dtype), mask=m_A) + + p_k = tl.advance(p_k, (K,)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_fwd_kernel_V( + q, + v, + z, + h, + o, + A, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_p = tl.maximum(i_t * BT - 1, 0) + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_z = tl.load(p_z, boundary_check=(0, 1)) + # [BT, BK] + b_zp = tl.load(p_zp, boundary_check=(0,)) + b_q = (b_q * exp(b_zp[None, :] - b_z)).to(b_q.dtype) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # works but dkw, owing to divine benevolence + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_q, b_h, allow_tf32=False) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_o += tl.dot(b_A.to(b_v.dtype), b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_dh( + q, + z, + do, + dh, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, + NORMK: tl.constexpr, +): + i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + b_zp = tl.full([BK if NORMK else BV], float('inf'), dtype=tl.float32) + for i_t in range(NT - 1, -1, -1): + i_p = tl.maximum(i_t * BT - 1, 0) + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BK, BT] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + if NORMK: + p_z = tl.make_block_ptr(z + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + b_r, b_zp = exp(b_zc - b_zp), b_zc + # [BK, BT] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_q = (b_q * exp(b_zc[:, None] - b_z)).to(b_q.dtype) + # [BK, BV] + b_dh = b_dh * b_r[:, None] + else: + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) + # [BV,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + b_r, b_zp = exp(b_zc - b_zp), b_zc + # [BT, BV] + b_z = tl.load(p_z, boundary_check=(0,)) + b_do = (b_do * exp(b_zc[None, :] - b_z)).to(b_do.dtype) + # [BK, BV] + b_dh = b_dh * b_r[None, :] + # [BK, BV] + b_dh += tl.dot(b_q, b_do, allow_tf32=False) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_V( + k, + v, + z, + h, + A, + do, + dh, + dq, + dk, + dv, + dA, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_p = tl.maximum(i_t * BT - 1, 0) + n_bh = tl.num_programs(2) + + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) + p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (0, i_t * BT), (BT, BT), (0, 1)) + + # [BK,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = exp(b_k - b_zc[None, :]).to(b_k.dtype) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * V * K, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BK, BV] + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + # [BT, BV] + b_dv = tl.dot(b_k, b_dh, allow_tf32=False) + if i_k == 0: + b_dv += tl.dot(b_A.to(b_do.dtype), b_do, allow_tf32=False) + b_do = (b_do * scale).to(b_do.dtype) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + # [BT, BT] + b_dA += tl.dot(b_do, tl.trans(b_v), allow_tf32=False) + # [BT, BK] + b_dq += tl.dot(b_do, b_h, allow_tf32=False) + # [BT, BK] + b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zp = tl.load(p_zp, boundary_check=(0,)) + # [BT, BK] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = exp(b_zp[None, :] - b_z) + # [BT, BK] + b_dq = b_dq * b_z + b_dk = b_dk * b_k + + p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + # [BT, BT] + b_dA = tl.where(m_s, b_dA, 0.).to(b_k.dtype) + if i_k == 0: + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_intra_V( + q, + k, + z, + dA, + dq, + dk, + T, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, +): + i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i = i_c // NC, i_c % NC + + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BK] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_zq = exp(b_zn[None, :] - b_z) + b_dq = tl.zeros([BC, BK], dtype=tl.float32) + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kz = exp(b_k - b_zn[None, :]).to(b_k.dtype) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + b_dq += tl.dot(b_dA, b_kz, allow_tf32=False) + b_dq *= b_zq + + o_i = tl.arange(0, BC) + o_dA = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC + m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, BC): + p_kj = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i*BC+j) * K + i_k * BK,), (BK,), (0,)) + # [BC,] + b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) + # [BK,] + b_kj = tl.load(p_kj, boundary_check=(0,)).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_dq += tl.where(m_i, b_dA[:, None] * exp(b_kj[None, :] - b_z), 0.) + p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*K, (T*K,), (1,), ((i_t * BT + i_i * BC + BC - 1) * K + i_k * BK,), (BK,), (0,)) + # [BK,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kz = exp(b_k - b_zn[None, :]) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_j * BC, i_i * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_qz = (b_q * exp(b_zn[None, :] - b_z)).to(b_q.dtype) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + b_dk += tl.dot(tl.trans(b_dA), b_qz, allow_tf32=False) + b_dk *= b_kz + + o_dA = i_bh * T * BT + (i_t * BT + i_i * BC) * BT + i_i * BC + tl.arange(0, BC) + for j in range(0, BC): + p_qj = tl.make_block_ptr(q + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) + p_zj = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) + # [BC,] + b_dA = tl.load(dA + o_dA + j * BT, mask=(i_t * BT + i_i * BC + j < T), other=0) + # [BK,] + b_qj = tl.load(p_qj, boundary_check=(0,)).to(tl.float32) + b_zj = tl.load(p_zj, boundary_check=(0,)).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp(b_k - b_zj[None, :]), 0.) + p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_intra_K( + v, + z, + do, + dA, + scale, + T, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + n_bh = tl.num_programs(2) + + if i_i > i_j: + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_t * BT + i_j * BC), (BV, BC), (0, 1)) + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_dA = tl.make_block_ptr(dA+(i_bh+i_v*n_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + # [BV,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BV] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_zn[None, :] - b_z) * scale).to(b_do.dtype) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = exp(b_v - b_zn[:, None]).to(b_v.dtype) + # [BC, BC] + b_dA = tl.dot(b_do, b_v, allow_tf32=False) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + elif i_i == i_j: + p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_j * BC) * V + i_v * BV,), (BV,), (0,)) + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + + o_i = tl.arange(0, BC) + o_A = (i_bh + i_v * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, BC): + # [BV,] + b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) + # [BC,] + b_dA = tl.sum(b_do * exp(b_v[None, :] - b_z), 1) + b_dA = tl.where(o_i >= j, b_dA, 0) + tl.store(dA + o_A + j, b_dA.to(b_do.dtype), mask=m_A) + + p_v = tl.advance(p_v, (V,)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_K( + q, + k, + v, + z, + h, + A, + do, + dh, + dq, + dk, + dv, + dA, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_p = tl.maximum(i_t * BT - 1, 0) + n_bh = tl.num_programs(2) + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh) * T * BT, (T, BT ), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k), allow_tf32=False) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) + p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) + p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + # [BV,] + b_zp = tl.load(p_zp, boundary_check=(0,)) + b_zc = tl.load(p_zc, boundary_check=(0,)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = exp(b_zp[None, :] - b_z) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * b_z * scale).to(b_do.dtype) + # [BK, BV] + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + # [BT, BK] + b_dq += tl.dot(b_do, b_h, allow_tf32=False) + b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) + # [BT, BV] + b_dv = b_v * tl.dot(b_k, b_dh, allow_tf32=False) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT ), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BT] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BT, BK] + b_dq += tl.dot(b_dA, b_k, allow_tf32=False) + b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, allow_tf32=False) + + p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_intra_KV( + v, + z, + A, + do, + dv, + T, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i = i_c // NC, i_c % NC + + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*V, (T*V,), (1,), ((i_t * BT + i_i * BC + BC - 1) * V + i_v * BV,), (BV,), (0,)) + # [BV,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_zn[None, :] - b_z)).to(b_do.dtype) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dv += tl.dot(b_A, b_do, allow_tf32=False) + b_dv *= exp(b_v - b_zn[None, :]) + + o_i = tl.arange(0, BC) + for j in range(0, BC): + p_z = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) + p_A = tl.make_block_ptr(A + i_bh * T * BT, (T * BT,), (1,), ((i_t * BT + i_i * BC + j) * BT + i_i * BC,), (BC,), (0,)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) + # [BC,] + b_A = tl.load(p_A, boundary_check=(0,)) + # [BV,] + b_z = tl.load(p_z, boundary_check=(0,)) + b_do = tl.load(p_do, boundary_check=(0,)) + # [BC, BV] + m_i = o_i[:, None] <= j + b_dv += tl.where(m_i, exp(b_v - b_z[None, :]) * b_A[:, None] * b_do[None, :], 0.) + p_dv = tl.make_block_ptr(dv + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_rcum_inter( + s, + z, + ss, + doo, + T, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + NT: tl.constexpr, +): + i_m, i_bh = tl.program_id(0), tl.program_id(1) + + b_sp = tl.zeros([BS], dtype=tl.float32) + b_zp = tl.full([BS], float('inf'), dtype=tl.float32) + for i_t in range(NT - 1, -1, -1): + p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) + p_zc = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT) * S + i_m * BS,), (BS,), (0,)) + p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) + p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) + # [BS,] + b_zc = tl.load(p_zc, boundary_check=(0,)) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)) + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_ss = tl.load(p_ss, boundary_check=(0, 1)) + + b_doo = exp(b_s - b_zp[None, :]) * b_sp[None, :] + tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) + # [BS,] + b_sp = b_sp * exp(b_zc - b_zp) + tl.sum(b_ss * exp(b_zc[None, :] - b_z), 0) + b_zp = b_zc + + +@triton.jit(do_not_specialize=['T']) +def chunk_abc_bwd_kernel_rcum_intra( + s, + z, + ss, + doo, + T, + S: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + NC: tl.constexpr, +): + i_s, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i = i_c // NC, i_c % NC + + o_i = tl.arange(0, BC) + m_o = tl.full([BC, BC], 1., dtype=tl.float32) + + p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) + p_zn = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + BC - 1) * S + i_s * BS,), (BS,), (0,)) + p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) + # [BC, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)) + # [BS,] + b_zn = tl.load(p_zn, boundary_check=(0,)) + + b_doo = tl.zeros([BC, BS], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) + p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) + # [BC, BS] + b_z = tl.load(p_z, boundary_check=(0, 1)) + b_ss = tl.load(p_ss, boundary_check=(0, 1)) + # [BC, BS] + b_doo += b_ss * exp(b_zn[None, :] - b_z) + b_doo = exp(b_s - b_zn[None, :]) * tl.dot(m_o.to(b_s.dtype), b_doo.to(b_s.dtype), allow_tf32=False) + + for j in range(0, BC): + p_z = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) + p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) + # [BS,] + b_z = tl.load(p_z, boundary_check=(0,)) + b_ss = tl.load(p_ss, boundary_check=(0,)) + # [BC, BS] + m_i = o_i[:, None] <= j + b_doo += tl.where(m_i, exp(b_s - b_z[None, :]) * b_ss[None, :], 0.) + b_doo += tl.load(p_doo, boundary_check=(0, 1)) + tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) + + +class ChunkABCFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward(ctx, q, k, v, s, initial_state, output_final_state): + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + BT, BC = 64, 16 + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + BM = min(64, triton.next_power_of_2(M)) + NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC) + NV, NM = triton.cdiv(V, BV), triton.cdiv(M, BM) + num_warps = 4 if BK == 64 else 2 + num_stages = 1 + + def fwd_pre(s, B, H, T, S): + # keep cummulative normalizer in fp32 + z = torch.empty_like(s, dtype=torch.float) + grid = (B * H,) + logcumsumexp_fwd_kernel[grid]( + s, z, + T=T, S=S, + ) + return z + + def fwd_inner(q, k, v, z, B, H, T, K, V, BT, BK, BV, NT, normk=False, h0=None, ht=None): + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + h = q.new_empty(B, H, NT * K, V) + grid = (NV, NK, B * H) + chunk_abc_fwd_kernel_h[grid]( + k, v, z, h, h0, ht, + T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT, + NORMK=normk, + USE_INITIAL_STATE=h0 is not None, + STORE_FINAL_STATE=ht is not None, + num_warps=num_warps, + num_stages=num_stages, + ) + return h + + final_state = None + if output_final_state: + final_state = (q.new_empty(B, H, K, M, dtype=torch.float), + q.new_empty(B, H, M, V, dtype=torch.float)) + + z = fwd_pre(s, B, H, T, M) + scale = K ** -0.5 + hk = fwd_inner( + q=q, k=k, v=s, z=z, + B=B, H=H, T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, + normk=False, + h0=initial_state[0] if initial_state is not None else None, + ht=final_state[0] if final_state is not None else None, + ) + ok1 = torch.empty_like(s) + Ak = q.new_empty(B, H, T, BT) + grid = (NM, NT, B * H) + chunk_abc_fwd_kernel_K[grid]( + q, k, z, hk, ok1, Ak, + scale=scale, + T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, + num_warps=num_warps, + num_stages=num_stages, + ) + ok0 = torch.empty_like(s) + grid = (NM, NT * NC, B * H) + chunk_abc_fwd_kernel_intra_K[grid]( + s, z, ok0, Ak, + T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, + num_warps=2, + num_stages=num_stages, + ) + ok = ok0.add_(ok1) + + scale = 1. + # p is kept in fp32 for safe softmax backward + p = softmax_fwd(ok, dtype=torch.float) + qv = p.to(q.dtype) + + scale = 1. + hv = fwd_inner( + q=qv, k=s, v=v, z=z, + B=B, H=H, T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, + normk=True, + h0=initial_state[1] if initial_state is not None else None, + ht=final_state[1] if final_state is not None else None, + ) + Av = q.new_zeros(NM, B, H, T, BT) + grid = (NM, NT * NC * NC, B * H) + chunk_abc_fwd_kernel_intra_V[grid]( + qv, s, z, Av, + scale=scale, + T=T, K=M, BT=BT, BC=BC, BK=BM, NC=NC, + num_warps=2, + num_stages=num_stages, + ) + Av = Av.sum(0) + ov = torch.empty_like(v) + grid = (NV, NT, B * H) + chunk_abc_fwd_kernel_V[grid]( + qv, v, z, hv, ov, Av, + scale=scale, + T=T, + K=M, + V=V, + BT=BT, + BK=BM, + BV=BV, + NT=NT, + num_warps=num_warps, + num_stages=num_stages, + ) + ctx.save_for_backward(q, k, v, s, z, ok, p, hk, hv, Av) + ctx.BT = BT + return ov, final_state + + @staticmethod + @input_guard + def backward(ctx, dov, dht=None): + q, k, v, s, z, ok, p, hk, hv, Av = ctx.saved_tensors + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + BT, BC = ctx.BT, 16 + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + BM = min(64, triton.next_power_of_2(M)) + NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC) + NK, NM = triton.cdiv(K, BK), triton.cdiv(M, BM) + num_warps = 4 if BK == 64 else 2 + num_stages = 1 + + def bwd_inner(q, z, do, B, H, T, K, V, BT, BK, BV, NT, scale, normk=False): + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + dh = q.new_empty(B, H, NT * K, V) + grid = (NK, NV, B * H) + chunk_abc_bwd_kernel_dh[grid]( + q, z, do, dh, + scale=scale, + T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT, + NORMK=normk, + num_warps=num_warps, + num_stages=num_stages, + ) + return dh + + def bwd_post(s, z, ss, B, H, T, S, BT, BC, BS, NT, NC, NS): + doo = torch.empty_like(s) + grid = (NS, B * H) + chunk_abc_bwd_kernel_rcum_inter[grid]( + s, z, ss, doo, + T=T, S=S, BT=BT, BS=BS, NT=NT, + num_warps=num_warps, + num_stages=num_stages, + ) + grid = (NS, NT * NC, B * H) + chunk_abc_bwd_kernel_rcum_intra[grid]( + s, z, ss, doo, + T=T, S=S, BT=BT, BC=BC, BS=BS, NC=NC, + num_warps=num_warps, + num_stages=num_stages, + ) + return doo + + scale = 1. + qv = p.to(q.dtype) + dhv = bwd_inner( + qv, z, dov, + B=B, H=H, T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, + scale=scale, + normk=True, + ) + dp1 = torch.empty_like(p) + dsv1 = torch.empty_like(s, dtype=torch.float) + dv = v.new_empty(NM, *v.shape) + dAv = q.new_zeros(B, H, T, BT) + grid = (NM, NT, B * H) + chunk_abc_bwd_kernel_V[grid]( + s, v, z, hv, Av, dov, dhv, dp1, dsv1, dv, dAv, + scale=scale, + T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, + num_warps=num_warps, + num_stages=num_stages, + ) + dv = dv.sum(0) + dp0 = torch.empty_like(p) + dsv0 = s.new_zeros(s.shape, dtype=torch.float) + grid = (NM, NT * NC, B * H) + chunk_abc_bwd_kernel_intra_V[grid]( + qv, s, z, dAv, dp0, dsv0, + T=T, K=M, BT=BT, BC=BC, BK=BM, NC=NC, + num_warps=2, + num_stages=num_stages, + ) + dp = dp1.add_(dp0) + dsv = dsv1.add_(dsv0) + + # softmax gradient, equivalent to: + # dok = p * (dp - (p * dp).sum(-1, True)) + dok = softmax_bwd(p, dp, dtype=ok.dtype) + + scale = K ** -0.5 + dhk = bwd_inner( + q, z, dok, + B=B, H=H, T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, + scale=scale, + normk=False, + ) + dAk = q.new_zeros(NM, B, H, T, BT) + grid = (NM, NT * NC * NC, B * H) + chunk_abc_bwd_kernel_intra_K[grid]( + s, z, dok, dAk, + scale=scale, + T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, + num_warps=2, + num_stages=num_stages, + ) + dAk = dAk.sum(0) + + Ak = q.new_zeros(NK, B, H, T, BT) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dsk1 = s.new_empty(NK, *s.shape, dtype=torch.float) + grid = (NK, NT, B * H) + chunk_abc_bwd_kernel_K[grid]( + q, k, s, z, hk, Ak, dok, dhk, dq, dk, dsk1, dAk, + scale=scale, + T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, + num_warps=num_warps, + num_stages=num_stages, + ) + Ak = Ak.sum(0) + dsk1 = dsk1.sum(0) + dsk0 = torch.empty_like(s, dtype=torch.float) + grid = (NM, NT * NC, B * H) + chunk_abc_bwd_kernel_intra_KV[grid]( + s, z, Ak, dok, dsk0, + T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, + num_warps=2, + num_stages=num_stages, + ) + ds = dsv.add_(dsk1.add_(dsk0)) + ds -= bwd_post(s, z, ok * dok + p * dp, B, H, T, M, BT, BC, BM, NT, NC, NM) + ds = ds.to(s.dtype) + return dq, dk, dv, ds, None, None + + +@torch.compiler.disable +def chunk_abc( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + initial_state: tuple[torch.Tensor] | None = None, + output_final_state: bool = False, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + s (torch.Tensor): + slot representations of shape `[B, T, H, M]`. + initial_state (Optional[Tuple[torch.Tensor, torch.Tensor]]): + Initial states of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `False`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[B, H, K, M]` and `[B, H, M, V]` if `output_final_state=True` else `None`. + """ + if not head_first: + q, k, v, s = map(lambda x: x.transpose(1, 2), (q, k, v, s)) + o, final_state = ChunkABCFunction.apply(q, k, v, s, initial_state, output_final_state) + if not head_first: + o = o.transpose(1, 2) + return o, final_state diff --git a/fla/ops/abc/naive.py b/fla/ops/abc/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..66f6193faf5e36026bc54c106a2f7defad613a47 --- /dev/null +++ b/fla/ops/abc/naive.py @@ -0,0 +1,94 @@ + + +import torch +from einops import repeat + + +def naive_recurrent_abc( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor | None = None, + scale: int | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, +) -> torch.Tensor: + dtype = q.dtype + + NG = q.shape[1]//k.shape[1] + # [batch_size, n_heads, seq_len, n_slots] + if g is None: + z = s.float().logcumsumexp(2) + g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z + s = torch.exp(s - z) + q, k, v, s, g = map(lambda x: x.float(), (q, k, v, s, g)) + k, v, s, g = map(lambda x: repeat(x, 'b h t d -> b (h g) t d', g=NG), (k, v, s, g)) + if initial_state is not None: + initial_state = tuple(map(lambda x: repeat(x, 'b h k v -> b (h g) k v', g=NG), initial_state)) + + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + + hk = torch.zeros(B, H, K, M, dtype=torch.float, device=q.device) + ok = torch.zeros_like(s) + + if scale is None: + scale = q.shape[-1] ** -0.5 + + final_state = None + if initial_state is not None: + hk += initial_state[0] + + for i in range(T): + q_i = q[:, :, i] * scale + k_i = k[:, :, i] + v_i = s[:, :, i] + g_i = g[:, :, i].exp() + hk = hk * g_i[..., None, :] + k_i[..., None] * v_i[..., None, :] + ok[:, :, i] = (q_i[..., None] * hk).sum(-2) + + qv = ok.softmax(-1) + hv = torch.zeros(B, H, M, V, dtype=torch.float, device=q.device) + ov = torch.zeros_like(v) + if initial_state is not None: + hv += initial_state[1] + + for i in range(T): + q_i = qv[:, :, i] + k_i = s[:, :, i] + v_i = v[:, :, i] + g_i = g[:, :, i].exp() + hv = hv * g_i[..., :, None] + k_i[..., None] * v_i[..., None, :] + ov[:, :, i] = (q_i[..., None] * hv).sum(-2) + + if output_final_state: + final_state = (hk.view(B, -1, NG, K, M)[:, :, 0], hv.view(B, -1, NG, M, V)[:, :, 0]) + return ov.to(dtype), final_state + + +def naive_cumsum_abc( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, +) -> torch.Tensor: + """ + A simple implementation of vanilla ABC that is more aligned with the descriptions in the paper. + This is just for demonstration purposes, with no numerical stabilities guaranteed. + """ + + dtype = q.dtype + q, k, v, s = map(lambda x: x.float(), (q, k, v, s)) + + scale = q.shape[-1] ** -0.5 + # [batch_size, n_heads, seq_len, n_slots] + s = (s - s.max(2, True)[0]).exp() + z = s.cumsum(2) + # [batch_size, n_heads, seq_len, n_slots, d_head] + K = (s.unsqueeze(-1) * k.unsqueeze(-2)).cumsum(2) / z.unsqueeze(-1) + V = (s.unsqueeze(-1) * v.unsqueeze(-2)).cumsum(2) / z.unsqueeze(-1) + # [batch_size, n_heads, seq_len, n_slots] + p = torch.einsum('...d,...md->...m', q * scale, K).softmax(-1) + # [batch_size, n_heads, seq_len, d_head] + o = torch.einsum('...m,...md->...d', p, V) + return o.to(dtype), None diff --git a/fla/ops/attn/__init__.py b/fla/ops/attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0b10a1e29f6763cdf06bf84fdb59386595700f9b --- /dev/null +++ b/fla/ops/attn/__init__.py @@ -0,0 +1,8 @@ + +from .naive import naive_parallel_attn +from .parallel import parallel_attn + +__all__ = [ + 'naive_parallel_attn', + 'parallel_attn', +] diff --git a/fla/ops/attn/decoding.py b/fla/ops/attn/decoding.py new file mode 100644 index 0000000000000000000000000000000000000000..e12e097beec327099316dc774bf0749aa86a3c74 --- /dev/null +++ b/fla/ops/attn/decoding.py @@ -0,0 +1,181 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.cumsum import chunk_global_cumsum +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, check_shared_mem + + +@triton.heuristics({ + 'USE_G': lambda args: args['g_cumsum'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4] + ([] if check_shared_mem('hopper') else [8]) + for num_stages in [2, 3, 4, 5] + ], + key=['H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'], + **autotune_cache_kwargs, +) +@triton.jit +def naive_attn_decoding_kernel( + q, + k, + v, + o, + g_cumsum, + scale, + gate_scale, + cu_seqlens, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, +): + i_v, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + bos, eos = tl.load(cu_seqlens + i_b).to(tl.int32), tl.load(cu_seqlens + i_b + 1).to(tl.int32) + T = eos - bos + + p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1, ), (0, ), (BK,), (0,)) + p_o = tl.make_block_ptr(o + i_bh * V, (V,), (1, ), (0, ), (BV,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0,)) + b_q = (b_q * scale).to(b_q.dtype) + + b_o = tl.zeros([BV ], dtype=tl.float32) + + b_m = tl.full([1], float('-inf'), dtype=tl.float32) + b_acc = tl.zeros([1], dtype=tl.float32) + + if USE_G: + p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (T-1,), (1,), (0,)) + b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + else: + b_gq = None + + for i_s in range(0, T, BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s, 0), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.sum(b_q[None, :] * b_k, 1) + + mask = i_s + tl.arange(0, BS) < T + b_s = tl.where(mask, b_s, float('-inf')) + + if USE_G: + p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + b_s += (b_gq - b_gk) * gate_scale + # [BT, BS] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s)), b_m + b_r = exp(b_mp - b_m) + # [BT, BS] + b_p = exp(b_s - b_m) + + # [BT] + b_acc = b_acc * b_r + tl.sum(b_p, 0) + # [BT, BV] + b_o = b_o * b_r + tl.sum(b_p[:, None] * b_v, 0) + b_mp = b_m + b_o = b_o / b_acc + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, )) + + +def attn_decoding_one_step( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor = None, + do_gate_scale: bool = False, +): + r""" + Args: + q (torch.Tensor): + query of shape `[1, B, HQ, K]`. + k (torch.Tensor): + keys of shape `[1, T, H, K]`. + GQA will be applied if HQ is divisible by H. T is the cumulative length for all batch. + v (torch.Tensor): + values of shape `[1, T, H, V]`. + g (Optional[torch.Tensor]): + log decay factors of shape `[1, T, H]`. Default: `None`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + do_gate_scale (bool): + Whether to apply gate scale. Default: `False`. If `True`, the attention scale will also be applied + to the gating bias term in Forgetting Transformer or PaTH-FoX. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, 1, HQ, V]`. + """ + assert cu_seqlens is not None, "The cu_seqlens must be provided for varlen decoding" + B, T, H, K, V = *k.shape, v.shape[-1] + N = len(cu_seqlens) - 1 + HQ = q.shape[2] + G = HQ // H + if scale is None: + scale = K ** -0.5 + + BK = max(triton.next_power_of_2(K), 16) + if check_shared_mem('hopper', q.device.index): + BS = min(64, max(16, triton.next_power_of_2(T))) + BV = min(256, max(16, triton.next_power_of_2(V))) + elif check_shared_mem('ampere', q.device.index): + BS = min(32, max(16, triton.next_power_of_2(T))) + BV = min(128, max(16, triton.next_power_of_2(V))) + else: + BS = min(32, max(16, triton.next_power_of_2(T))) + BV = min(64, max(16, triton.next_power_of_2(V))) + g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, output_dtype=torch.float32) if g is not None else None + NV = triton.cdiv(V, BV) + o = torch.empty(*q.shape[:-1], V, dtype=v.dtype, device=q.device) + gate_scale = 1.0 if not do_gate_scale else scale + + grid = (NV, N * HQ) + naive_attn_decoding_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + g_cumsum=g_cumsum, + scale=scale, + gate_scale=gate_scale, + cu_seqlens=cu_seqlens, + B=B, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BS=BS, + BK=BK, + BV=BV, + ) + return o diff --git a/fla/ops/attn/naive.py b/fla/ops/attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..75e5a1a6ec3e645ff73098689f349467c0dbb184 --- /dev/null +++ b/fla/ops/attn/naive.py @@ -0,0 +1,64 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F + + +def naive_parallel_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + causal: bool = True, +): + """ + Reference PyTorch implementation of parallel attention that returns both output and max_logits. + + Args: + q: [B, T, HQ, D] + k: [B, T, H, D] + v: [B, T, H, D] + scale: float, optional. If None, defaults to 1 / sqrt(D) + causal: bool, default True + + Returns: + output: [B, T, HQ, D] + max_logits: [B, T, HQ] + """ + B, T, HQ, D = q.shape + H = k.shape[2] + G = HQ // H + + if scale is None: + scale = D ** -0.5 + + # Reshape q to separate heads and groups + q = q.reshape(B, T, H, G, D) # [B, T, H, G, D] + + # Repeat k and v to match groups: [B, T, H, D] -> [B, T, H, G, D] + k = k.unsqueeze(3).expand(B, T, H, G, D) # Expand along group dimension + v = v.unsqueeze(3).expand(B, T, H, G, D) + + # Reshape to treat each (B, H, G,) as a separate head + q_flat = q.reshape(B * H * G, T, D) # [B*H*G, T, D] + k_flat = k.reshape(B * H * G, T, D) # [B*H*G, T, D] + v_flat = v.reshape(B * H * G, T, D) # [B*H*G, T, D] + + # Compute attention scores: [B*H*G, T, T] + scores = torch.bmm(q_flat, k_flat.transpose(1, 2)) * scale + + # Apply causal mask + if causal: + causal_mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=q.device), diagonal=1) + scores = scores.masked_fill(causal_mask.unsqueeze(0), float('-inf')) + + # Compute max_logits (max over key dimension): [B*H*G, T] + max_logits_flat = scores.max(dim=-1).values + max_logits = max_logits_flat.reshape(B, T, HQ) # [B, T, HQ] + + # Compute attention weights and output + attn_weights = F.softmax(scores, dim=-1) # [B*H*G, T, T] + output_flat = torch.bmm(attn_weights, v_flat) # [B*H*G, T, D] + output = output_flat.reshape(B, T, HQ, D) # [B, T, HQ, D] + + return output, max_logits diff --git a/fla/ops/attn/parallel.py b/fla/ops/attn/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..b6b73ffcf72f345d6bd6b07f3f8cb839bef3de42 --- /dev/null +++ b/fla/ops/attn/parallel.py @@ -0,0 +1,734 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl +from einops import reduce + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cumsum import chunk_global_cumsum +from fla.ops.utils.op import exp2, log2 +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous + + +@triton.heuristics({ + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit +def parallel_attn_fwd_kernel( + q, + k, + v, + o, + g_cumsum, + lse, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + RCP_LN2: tl.constexpr = 1.4426950216 + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BT, BV] + b_o = tl.zeros([BT, BV], dtype=tl.float32) + + b_m = tl.full([BT], float('-inf'), dtype=tl.float32) + b_acc = tl.zeros([BT], dtype=tl.float32) + + if USE_G: + p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + else: + b_gq = None + + for i_s in range(0, i_t * BT, BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + + if USE_G: + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) + b_s += b_gq[:, None] - b_gk[None, :] + + # [BT, BS] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp2(b_mp - b_m) + # [BT, BS] + b_p = exp2(b_s - b_m[:, None]) + # [BT] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [BT, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + + # [BS] + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + + if USE_G: + b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) + b_s += b_gq[:, None] - b_gk[None, :] + + b_s = tl.where((o_q[:, None] >= o_k[None, :]) & m_k[None, :], b_s, float('-inf')) + + # [BT] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp2(b_mp - b_m) + # [BT, BS] + b_p = exp2(b_s - b_m[:, None]) + # [BT] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [BT, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + b_mp = b_m + + b_o = b_o / b_acc[:, None] + b_m += log2(b_acc) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) + + +@triton.jit +def parallel_attn_bwd_kernel_preprocess( + o, + do, + delta, + B: tl.constexpr, + V: tl.constexpr, +): + i_n = tl.program_id(0) + o_d = tl.arange(0, B) + m_d = o_d < V + + b_o = tl.load(o + i_n * V + o_d, mask=m_d, other=0) + b_do = tl.load(do + i_n * V + o_d, mask=m_d, other=0).to(tl.float32) + b_delta = tl.sum(b_o * b_do) + + tl.store(delta + i_n, b_delta.to(delta.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_attn_bwd_kernel_dq( + q, + k, + v, + lse, + delta, + do, + dq, + dg_cumsum, + g_cumsum, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + # NOTE: we must multiply RCP_LN2 after tl.dot for high precision + RCP_LN2: tl.constexpr = 1.4426950216 + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT] + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + + # [BT, BK] + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + if USE_G: + b_dg = tl.zeros([BT], dtype=tl.float32) + p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + else: + b_gq = None + b_dg = None + + o_q = i_t * BT + tl.arange(0, BT) + for i_s in range(0, i_t * BT, BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + if USE_G: + b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) + b_s += b_gq[:, None] - b_gk[None, :] + + b_s = tl.where((o_q[:, None] >= o_k[None, :]) & m_k[None, :], b_s, float('-inf')) + b_p = exp2(b_s - b_lse[:, None]) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + if USE_G: + b_dg += tl.sum(b_ds, 1) + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + + # [BS] + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + + if USE_G: + p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + b_s += b_gq[:, None] - b_gk[None, :] + b_p = tl.where((o_q[:, None] >= o_k[None, :]) & m_k[None, :], exp2(b_s - b_lse[:, None]), 0) + + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + if USE_G: + b_dg += tl.sum(b_ds, 1) + + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + if USE_G: + p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_attn_bwd_kernel_dkv( + q, + k, + v, + g_cumsum, + lse, + delta, + do, + dk, + dv, + dg_cumsum, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + RCP_LN2: tl.constexpr = 1.4426950216 + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + o_k = i_t * BT + tl.arange(0, BT) + + if USE_G: + p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + b_dg = tl.zeros([BT], dtype=tl.float32) + else: + b_gk = None + b_dg = None + + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_q = i_s + tl.arange(0, BS) + m_q = o_q < T + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BS] + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + # [BT, BS] + b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 + if USE_G: + p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + b_s += b_gq[None, :] - b_gk[:, None] + b_p = tl.where((o_k[:, None] <= o_q[None, :]) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) + # [BT, BS] @ [BS, BV] -> [BT, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BT, BS] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + if USE_G: + b_dg -= tl.sum(b_ds, 1) + + for i_s in range((i_t + 1) * BT, tl.cdiv(T, BS) * BS, BS): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_q = i_s + tl.arange(0, BS) + m_q = o_q < T + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BS] + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + # [BT, BS] + b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 + if USE_G: + p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + b_s += b_gq[None, :] - b_gk[:, None] + b_p = tl.where(m_q[None, :], exp2(b_s - b_lse[None, :]), 0) + # [BT, BS] @ [BS, BV] -> [BT, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BT, BS] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + if USE_G: + b_dg -= tl.sum(b_ds, 1) + + b_dk = b_dk * scale + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + if USE_G: + p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +def parallel_attn_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + BT = 128 + if check_shared_mem('hopper', q.device.index): + BS = min(64, max(16, triton.next_power_of_2(T))) + BK = min(256, max(16, triton.next_power_of_2(K))) + BV = min(256, max(16, triton.next_power_of_2(V))) + num_warps = 8 + elif check_shared_mem('ampere', q.device.index): + BS = min(32, max(16, triton.next_power_of_2(T))) + BK = min(256, max(16, triton.next_power_of_2(K))) + BV = min(128, max(16, triton.next_power_of_2(V))) + num_warps = 4 + else: + BS = min(32, max(16, triton.next_power_of_2(T))) + BK = min(256, max(16, triton.next_power_of_2(K))) + BV = min(64, max(16, triton.next_power_of_2(V))) + num_warps = 2 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert NK == 1, "The key dimension can not be larger than 256" + + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + grid = (NV, NT, B * HQ) + parallel_attn_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + g_cumsum=g_cumsum, + lse=lse, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + num_warps=num_warps, + ) + return o, lse + + +def parallel_attn_bwd_preprocess( + o: torch.Tensor, + do: torch.Tensor, +): + V = o.shape[-1] + delta = torch.empty_like(o[..., 0], dtype=torch.float) + parallel_attn_bwd_kernel_preprocess[(delta.numel(),)]( + o=o, + do=do, + delta=delta, + B=triton.next_power_of_2(V), + V=V, + ) + return delta + + +def parallel_attn_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + g_cumsum: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + scale: float = None, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + if check_shared_mem('hopper'): + BT = 128 + BS = 64 + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + num_warps = 8 + elif check_shared_mem('ampere'): + BS = 32 + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + BT = 128 if K <= 64 else 64 + num_warps = 4 + else: + BT = 64 + BS = 32 + BK = max(triton.next_power_of_2(K), 16) + BV = min(max(triton.next_power_of_2(V), 16), 64) + num_warps = 2 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NV = triton.cdiv(V, BV) + + delta = parallel_attn_bwd_preprocess(o, do) + + dq = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) + dk = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) + dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) + grid = (NV, NT, B * HQ) + + dg_cumsum, dg_cumsum_k = None, None + if g_cumsum is not None: + dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + parallel_attn_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + lse=lse, + delta=delta, + do=do, + dq=dq, + dg_cumsum=dg_cumsum, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + num_warps=num_warps, + ) + parallel_attn_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + dg_cumsum=dg_cumsum_k, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + num_warps=num_warps, + ) + dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') + dv = reduce(dv, 'b t (h g) v -> b t h v', g=G, reduction='sum') + if g_cumsum is not None: + dg_cumsum.add_(dg_cumsum_k) + return dq, dk, dv, dg_cumsum + + +@torch.compile +class ParallelAttentionFunction(torch.autograd.Function): + + @staticmethod + @contiguous + @autocast_custom_fwd + def forward(ctx, q, k, v, g, scale, cu_seqlens, chunk_indices=None): + ctx.dtype = q.dtype + + RCP_LN2: float = 1.4426950216 + g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, scale=RCP_LN2) if g is not None else None + o, lse = parallel_attn_fwd( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, k, v, o, g_cumsum, lse) + ctx.cu_seqlens = cu_seqlens + ctx.scale = scale + return o.to(q.dtype) + + @staticmethod + @contiguous + @autocast_custom_bwd + def backward(ctx, do): + q, k, v, o, g_cumsum, lse = ctx.saved_tensors + dq, dk, dv, dg = parallel_attn_bwd( + q=q, + k=k, + v=v, + o=o, + g_cumsum=g_cumsum, + lse=lse, + do=do, + scale=ctx.scale, + cu_seqlens=ctx.cu_seqlens, + ) + if dg is not None: + dg = chunk_global_cumsum(dg, cu_seqlens=ctx.cu_seqlens, reverse=True) + + return dq.to(q), dk.to(k), dv.to(v), dg, None, None, None + + +def parallel_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + GQA will be applied if HQ is divisible by H. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (Optional[torch.Tensor]): + log decay factors of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + + o = ParallelAttentionFunction.apply(q, k, v, g, scale, cu_seqlens, chunk_indices) + return o diff --git a/fla/ops/backends/__init__.py b/fla/ops/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..611158a5a7a6a071142fd3bc02ae9ea5d494d873 --- /dev/null +++ b/fla/ops/backends/__init__.py @@ -0,0 +1,179 @@ +"""Generic backend dispatch system for FLA operations.""" + +from __future__ import annotations + +import contextlib +import logging +import os +import threading +from collections.abc import Callable +from functools import wraps +from importlib.util import find_spec +from typing import Any, ClassVar, TypeVar + +logger = logging.getLogger(__name__) +F = TypeVar('F', bound=Callable) + + +class BaseBackend: + """Base class for operation-specific backends. + + Attributes: + backend_type: Identifier for the backend type, used to distinguish different backend implementations. + package_name: Name of the external package required by the backend. None indicates no external dependency. + env_var: Environment variable name that controls whether the backend is enabled. None means always enabled. + default_enable: Controls whether the backend is enabled by default when env_var is not set. + Defaults to True (enabled). Set to False to require explicit user opt-in. + priority: Backend priority, lower values indicate higher priority (default is 5). + """ + + backend_type: ClassVar[str] = "base" + package_name: ClassVar[str | None] = None + env_var: ClassVar[str | None] = None + default_enable: ClassVar[bool] = True + # Lower number = higher priority, default is 5 + priority: ClassVar[int] = 5 + + @classmethod + def is_available(cls) -> bool: + if cls.package_name is None: + return True + return find_spec(cls.package_name) is not None + + @classmethod + def is_enabled(cls) -> bool: + if cls.env_var is None: + return True + default_value = "1" if cls.default_enable else "0" + return os.environ.get(cls.env_var, default_value) != "0" + + @classmethod + def can_use(cls) -> bool: + return cls.is_available() and cls.is_enabled() + + def verify(self, func_name: str, *args, **kwargs) -> tuple[bool, str | None]: + """Check if backend can handle the function call.""" + verifier_name = f"{func_name}_verifier" + verifier = getattr(self, verifier_name, None) + if verifier is None: + return True, None + + try: + return verifier(*args, **kwargs) + except Exception as e: + return False, str(e) + + +class BackendRegistry: + """Per-operation backend registry.""" + + _registries: ClassVar[dict[str, BackendRegistry]] = {} + _initialized: ClassVar[set[str]] = set() + _init_lock: ClassVar[threading.Lock] = threading.Lock() + + def __init__(self, operation_name: str): + self.operation_name = operation_name + self._backends: dict[str, BaseBackend] = {} + self._active: BaseBackend | None = None + self._lock = threading.RLock() + self._logged: set[str] = set() + BackendRegistry._registries[operation_name] = self + + def register(self, backend: BaseBackend) -> None: + """Register a backend.""" + with self._lock: + self._backends[backend.backend_type] = backend + # Update active backend based on priority + self._update_active_backend() + + def _get_sorted_backends(self) -> list[BaseBackend]: + """Get backends sorted by priority (lower number = higher priority). + + Backends with the same priority are sorted by registration order. + """ + return sorted( + self._backends.values(), + key=lambda b: (b.priority, list(self._backends.values()).index(b)) + ) + + def _update_active_backend(self) -> None: + """Update active backend based on priority.""" + for backend in self._get_sorted_backends(): + if backend.can_use(): + self._active = backend + return + + def get_active(self) -> BaseBackend | None: + """Get active backend.""" + return self._active + + @classmethod + def ensure_initialized(cls, operation: str) -> None: + """Lazy-load backends on first use.""" + if operation in cls._initialized: + return + + with cls._init_lock: + if operation in cls._initialized: + return + + # Import backend module to trigger registration + with contextlib.suppress(ImportError): + __import__(f'fla.ops.{operation}.backends', fromlist=['']) + + cls._initialized.add(operation) + + +def dispatch(operation: str): + """Dispatch decorator with verifier support. + + Iterates through all registered backends and selects the first one + that passes the verifier for the given function call. + """ + def decorator(func: F) -> F: + func_name = func.__name__ + + @wraps(func) + def wrapper(*args, **kwargs) -> Any: + # Lazy initialization of backends + BackendRegistry.ensure_initialized(operation) + + registry = BackendRegistry._registries.get(operation) + if registry is None: + return func(*args, **kwargs) + + # Iterate through all registered backends sorted by priority + # to find one that can handle this call + backends_list = registry._get_sorted_backends() + + for be in backends_list: + if not be.can_use(): + continue + + can_use, _ = be.verify(func_name, *args, **kwargs) + if not can_use: + continue + + impl = getattr(be, func_name, None) + if impl is None: + continue + + result = impl(*args, **kwargs) + + log_key = f"{operation}:{func_name}:{be.backend_type}" + if log_key not in registry._logged: + with registry._lock: + if log_key not in registry._logged: + registry._logged.add(log_key) + logger.info(f"[FLA Backend] {operation}.{func_name} -> {be.backend_type}") + + return result + + # No backend can handle this call, use default implementation + return func(*args, **kwargs) + + return wrapper + return decorator + + +__all__ = ['BackendRegistry', 'BaseBackend', 'dispatch'] diff --git a/fla/ops/based/__init__.py b/fla/ops/based/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d73959049e530e75fadb63ba7633bb425d045b3c --- /dev/null +++ b/fla/ops/based/__init__.py @@ -0,0 +1,8 @@ + +from .fused_chunk import fused_chunk_based +from .parallel import parallel_based + +__all__ = [ + 'fused_chunk_based', + 'parallel_based', +] diff --git a/fla/ops/based/fused_chunk.py b/fla/ops/based/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..2ad853fa091ab5da19b025de56ecbd0f49ba2e36 --- /dev/null +++ b/fla/ops/based/fused_chunk.py @@ -0,0 +1,371 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.jit(do_not_specialize=['T']) +def fused_chunk_based_fwd_kernel( + q, + k, + v, + o, + z, + scale, # K ** -0.5 + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + o_i = tl.arange(0, BT) + + # [BT, BT] + m_s = o_i[:, None] >= o_i[None, :] + + # [BV], zero-order taylor expansion + b_h_0o = tl.zeros([BV], dtype=tl.float32) + # [BK, BV], first-order taylor expansion + b_h_1o = tl.zeros([BK, BV], dtype=tl.float32) + # [BK, BK, BV] second-order taylor expansion + b_h_2o = tl.zeros([BK*BK, BV], dtype=tl.float32) + + # make block pointers + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (i_bh + i_k*B*H) * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) + + p_z = z + (i_bh + i_k * B * H) * T + tl.arange(0, BT) + k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) + k_1o = tl.zeros([1, BK], dtype=tl.float32) + k_0o = 0 + + for i in range(0, tl.cdiv(T, BT)): + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK*BK, BT] + b_k_2o = b_k[:, None, :] * b_k[None, :, :] + b_k_2o = tl.reshape(b_k_2o, [BK * BK, BT]).to(b_k.dtype) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BK] + b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(b_k.dtype) + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_z = tl.zeros([BT], dtype=tl.float32) + + # interchunk + # zero-order + b_o += b_h_0o + b_z += k_0o + # first-order + b_o += tl.dot(b_q, b_h_1o.to(b_q.dtype), allow_tf32=False) + b_z += tl.sum(b_q * k_1o, axis=1) + # second-order + b_q_2o = b_q[:, :, None] * b_q[:, None, :] + b_q_2o = tl.reshape(b_q_2o, [BT, BK * BK]).to(b_k.dtype) + b_o += tl.dot(b_q_2o, b_h_2o.to(b_q_2o.dtype), allow_tf32=False) * 0.5 + b_z += tl.sum(b_q_2o * k_2o, axis=1) * 0.5 + + # update running statistics + k_1o += tl.sum(b_k, axis=1)[None, :] + k_2o += tl.sum(b_k_2o, axis=1)[None, :] + k_0o += BT + + # intrachunk + # [BT, BT] + b_s = tl.dot(b_q, b_k, allow_tf32=False) + b_s = 1 + b_s + 0.5 * b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_z += tl.sum(b_s, axis=1) + b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + # [TB, BV] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=(i * BT + tl.arange(0, BT)) < T) + + # update hidden state + # [BK, BV] + b_h_2o = b_h_2o + tl.dot(b_k_2o.to(b_v.dtype), b_v, allow_tf32=False) + b_h_1o = b_h_1o + tl.dot(b_k, b_v, allow_tf32=False) + b_h_0o = b_h_0o + tl.sum(b_v, axis=0) + + p_q = tl.advance(p_q, (BT, 0)) + p_k = tl.advance(p_k, (0, BT)) + p_v = tl.advance(p_v, (BT, 0)) + p_o = tl.advance(p_o, (BT, 0)) + p_z += BT + + +# Similar to Algorithm1 of https://arxiv.org/abs/2006.16236 +@triton.jit +def fused_chunk_based_bwd_kernel( + # NV: number of split in the V dimension. NK: number of split in the K dimension + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, # K ** -0.5 + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + # [BV], zero-order taylor expansion + # b_h_0o = tl.zeros([BV], dtype=tl.float32) + # [BK, BV], first-order taylor expansion + b_h_1o = tl.zeros([BV, BK], dtype=tl.float32) + # [BK, BK, BV] second-order taylor expansion + b_h_2o = tl.zeros([BV, BK*BK], dtype=tl.float32) + + k_1o = tl.zeros([1, BK], dtype=tl.float32) + k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) + + for i in range(0, tl.cdiv(T, BT)): + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i * BT, i_v * BV), (BT, BV), (1, 0)) + p_dq = tl.make_block_ptr(dq + (i_bh + i_v*B*H) * T*K, (T, K), (K, 1), (i*BT, i_k*BK), (BT, BK), (1, 0)) + p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i * BT + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + + # load tensors + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(tl.arange(0, BT) + i * BT) < T) + # [BV, BT] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # inter-chunk + b_dq += tl.dot(b_do, (b_h_1o).to(b_do.dtype), allow_tf32=False) + if i_v == 0: + b_dq += b_dz[:, None] * k_1o + b_dq_2o = tl.dot(b_do, (b_h_2o).to(b_do.dtype), allow_tf32=False) * 0.5 + if i_v == 0: + b_dq_2o += (b_dz[:, None] * k_2o) * 0.5 + b_dq_2o = tl.reshape(b_dq_2o, [BT, BK, BK]) + b_dq += tl.sum(b_dq_2o * b_q[:, :, None], axis=1) + b_dq += tl.sum(b_dq_2o * b_q[:, None, :], axis=2) + b_dq *= scale + + # intra-chunk + # [BT, BT] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + b_ds = tl.where(m_s, b_ds, 0) * scale + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + b_s = tl.where(m_s, b_s, 0) + b_dq += tl.dot((b_ds * (1 + b_s)).to(b_q.dtype), b_k, allow_tf32=False) + + # store + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + # update hidden state + # [BT, BK*BK] + b_k_2o = b_k[:, :, None] * b_k[:, None, :] + b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) + # [BV, BK*BK] + b_h_2o = b_h_2o + tl.dot(b_v, b_k_2o.to(b_v.dtype), allow_tf32=False) + # [BV, BK] + b_h_1o = b_h_1o + tl.dot(b_v, b_k, allow_tf32=False) + + if i_v == 0: + # update running statistics + k_1o += tl.sum(b_k, axis=0)[None, :] + k_2o += tl.sum(b_k_2o, axis=0)[None, :] + + tl.debug_barrier() + b_h_1o = None + b_h_2o = None + + # [BK, BV], first-order taylor expansion + b_dh_1o = tl.zeros([BK, BV], dtype=tl.float32) + # [BK, BK, BV] second-order taylor expansion + b_dh_2o = tl.zeros([BK*BK, BV], dtype=tl.float32) + b_dh_0o = tl.zeros([BV], dtype=tl.float32) + m_s = tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :] + + dq_1o = tl.zeros([1, BK], dtype=tl.float32) + dq_2o = tl.zeros([BK * BK, 1], dtype=tl.float32) + + for i in range(tl.cdiv(T, BT) * BT - BT, -BT, -BT): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_bh+i_v*B*H) * T*K, (T, K), (K, 1), (i, i_k*BK), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_bh+i_k*B*H) * T*V, (T, V), (V, 1), (i, i_v*BV), (BT, BV), (1, 0)) + p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i + + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(tl.arange(0, BT)+i) < T) + b_q = (b_q * scale).to(b_k.dtype) + + # intra chunk + b_ds = tl.dot(b_v, tl.trans(b_do), allow_tf32=False) + if i_v == 0: + b_ds += b_dz[None, :] + b_ds = tl.where(m_s, b_ds, 0) + b_s = tl.dot(b_k, b_q, allow_tf32=False) + b_s2 = 1 + b_s + 0.5 * b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_s2 = tl.where(m_s, b_s2, 0) + b_ds *= (1+b_s) + + b_dk += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q), allow_tf32=False) + b_dv += tl.dot(b_s2.to(b_do.dtype), b_do, allow_tf32=False) + + # inter chunk + b_k_2o = b_k[:, :, None] * b_k[:, None, :] + b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) + + b_dv += tl.dot(b_k, b_dh_1o.to(b_k.dtype), allow_tf32=False) + b_dv += tl.dot(b_k_2o, b_dh_2o.to(b_k.dtype), allow_tf32=False) + b_dv += b_dh_0o + + b_dk += tl.dot(b_v, tl.trans(b_dh_1o).to(b_k.dtype), allow_tf32=False) + + if i_v == 0: + b_dk += dq_1o + + b_dk_2o = tl.dot(b_dh_2o.to(b_k.dtype), tl.trans(b_v), allow_tf32=False) + if i_v == 0: + b_dk_2o += dq_2o + b_dk_2o = tl.reshape(b_dk_2o, [BK, BK, BT]) + b_k_fp32 = tl.trans(b_k.to(tl.float32)) + b_dk2 = tl.sum(b_dk_2o * b_k_fp32[:, None, :], axis=0) + b_dk2 += tl.sum(b_dk_2o * b_k_fp32[None, :, :], axis=1) + b_dk += tl.trans(b_dk2) + + # hidden state update + b_dh_0o += tl.sum(b_do, axis=0) + b_dh_1o = b_dh_1o + tl.dot(b_q, b_do, allow_tf32=False) + b_q_2o = b_q[None, :, :] * b_q[:, None, :] + b_q_2o = tl.reshape(b_q_2o, [BK * BK, BT]).to(b_k.dtype) + b_dh_2o = b_dh_2o + tl.dot(b_q_2o, b_do, allow_tf32=False) * 0.5 + + if i_v == 0: + dq_1o += (tl.sum(b_dz[None, :] * b_q, axis=1))[None, :] + dq_2o += (tl.sum(b_dz[None, :] * b_q_2o, axis=1) * 0.5)[:, None] + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +class FusedChunkBasedFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale=1): + B, H, T, K, V = *k.shape, v.shape[-1] + + scale = scale + BT = 16 + BK, BV = min(K, 16), min(V, 32) + BK, BV = max(BK, 16), max(BV, 16) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + num_warps = 4 + + # the norm of o might explode, so we need to use float32 here + o = q.new_empty(NK, B, H, T, V, dtype=torch.float32) + z = q.new_empty(NK, B, H, T, dtype=torch.float32) + + grid = (NV, NK, B * H) + fused_chunk_based_fwd_kernel[grid]( + q, k, v, o, z, + scale, + T=T, B=B, H=H, K=K, V=V, BT=BT, BK=BK, BV=BV, + num_warps=num_warps, + ) + o = o.sum(0) + z = z.sum(0) + ctx.save_for_backward(q, k, v) + ctx.scale = scale + return o.to(q.dtype), z.to(z.dtype) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dz): + q, k, v = ctx.saved_tensors + B, H, T, K, V = *k.shape, v.shape[-1] + scale = ctx.scale + + BT = 16 + BK, BV = min(K, 16), min(V, 32) + BK, BV = max(BK, 16), max(BV, 16) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + num_stages = 1 + num_warps = 4 + + dq = q.new_empty(NV, B, H, T, K) + dk = q.new_empty(NV, B, H, T, K) + dv = q.new_empty(NK, B, H, T, V) + grid = (NV, NK, B * H) + + fused_chunk_based_bwd_kernel[grid]( + q, k, v, do, dz, dq, dk, dv, + scale, + T=T, B=B, H=H, K=K, V=V, BT=BT, BK=BK, BV=BV, + num_warps=num_warps, + num_stages=num_stages, + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), None + + +def fused_chunk_based( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + use_norm: bool = True, + head_first: bool = False, +): + assert q.shape[-1] <= 16, 'only support feature dimension up to 16.' + if scale is None: + scale = q.shape[-1] ** -0.5 + if not head_first: + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + o, z = FusedChunkBasedFunction.apply(q, k, v, scale) + if use_norm: + o = o / (z[..., None] + 1e-6) + if not head_first: + o = o.transpose(1, 2) + return o.to(q.dtype) diff --git a/fla/ops/based/naive.py b/fla/ops/based/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..cdc8f4ffae572050385adada1ceb11101d4e88e1 --- /dev/null +++ b/fla/ops/based/naive.py @@ -0,0 +1,70 @@ + + +import torch +from einops import rearrange + + +def naive_parallel_based( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + use_norm: bool = True, +): + if scale is None: + scale = q.shape[-1] ** -0.5 + q = q * scale + attn = q @ k.transpose(-2, -1) + attn = 1 + attn + 1/2 * (attn ** 2) + attn.masked_fill_(~torch.tril(torch.ones( + q.shape[-2], q.shape[-2], dtype=torch.bool, device=q.device)), 0) + o = attn @ v + if use_norm: + z = attn.sum(-1) + return o / (z[..., None] + 1e-6) + else: + return o + + +def naive_chunk_based(q, k, v, chunk_size=256): + q = q * (q.shape[-1] ** -0.5) + # compute normalizer. + k_cumsum = torch.cumsum(k, dim=-2) + kk_cumsum = torch.cumsum(k.unsqueeze(-1) * k.unsqueeze(-2), dim=-3) + # first + z = (q * k_cumsum).sum(-1) + # second order + z += (q.unsqueeze(-1) * q.unsqueeze(-2) * kk_cumsum).sum((-1, -2)) * 0.5 + # zero-th order + z += (torch.arange(0, q.shape[-2]).to(z.device) * 1.0 + 1.0)[None, None, :] + + # compute o + # constant term + _o = v.cumsum(-2) + + q = rearrange(q, 'b h (n c) d -> b h n c d', c=chunk_size) + + k = rearrange(k, 'b h (n c) d -> b h n c d', c=chunk_size) + v = rearrange(v, 'b h (n c) d -> b h n c d', c=chunk_size) + + intra_chunk_attn = q @ k.transpose(-2, -1) + intra_chunk_attn = intra_chunk_attn + 1/2 * (intra_chunk_attn ** 2) + intra_chunk_attn.masked_fill_(~torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device)), 0) + o = intra_chunk_attn @ v + + # quadractic term + kv = torch.einsum('b h n c x, b h n c y, b h n c z -> b h n x y z', k, k, v) + kv = kv.cumsum(2) + kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) + + o += 0.5 * torch.einsum('b h n x y z, b h n c x, b h n c y -> b h n c z', kv, q, q) + + # linear term + kv = torch.einsum('b h n c x, b h n c y -> b h n x y', k, v) + kv = kv.cumsum(2) + kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) + o += torch.einsum('b h n x y, b h n c x -> b h n c y', kv, q) + + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o + _o + return o / (z[..., None] + 1e-6) diff --git a/fla/ops/based/parallel.py b/fla/ops/based/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..71f10f53f8bd45ce33b5c61b087afe6e81e96d37 --- /dev/null +++ b/fla/ops/based/parallel.py @@ -0,0 +1,406 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + +# Based: An Educational and Effective Sequence Mixer +# https://hazyresearch.stanford.edu/blog/2023-12-11-zoology2-based + + +@triton.jit(do_not_specialize=['T']) +def parallel_based_fwd_kernel( + q, + k, + v, + o, + z, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # i_c: chunk index. used for sequence parallelism + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % (NV) + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BTS, BV), (1, 0)) + + # [BQ, BD] block Q, in the shared memory throughout the whole kernel + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + b_o = tl.zeros([BTL, BV], dtype=tl.float32) + b_z = tl.zeros([BTL], dtype=tl.float32) + + # Q block and K block have no overlap + # no need for mask, thereby saving flops + for _ in range(0, i_c * BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_s = tl.dot(b_q, (b_k), allow_tf32=False) + b_s = 1 + b_s + 0.5 * b_s * b_s + b_z += tl.sum(b_s, axis=1) + + # [BQ, BD] + b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + + # # rescale interchunk output + tl.debug_barrier() + o_q = tl.arange(0, BTL) + # # sync threads, easy for compiler to optimize + # tl.debug_barrier() + + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_c * BTL), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTS, BV), (1, 0)) + # Q block and K block have overlap. masks required + for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_s = tl.dot(b_q, b_k, allow_tf32=False) + b_s = 1 + b_s + 0.5 * b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_z += tl.sum(b_s, axis=1) + # [BTL, BV] + b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + o_k += BTS + + p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_z = z + (i_bh + B * H * i_k) * T + i_c * BTL + tl.arange(0, BTL) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c * BTL + tl.arange(0, BTL)) < T)) + + +@triton.jit +def _parallel_based_bwd_dq( + i_bh, + i_c, + i_k, + i_v, + q, + k, + v, + do, + dz, + dq, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, +): + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) + p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dq = tl.zeros([BTL, BK], dtype=tl.float32) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, 0), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i_c * BTL + tl.arange(0, BTL) + b_dz = tl.load(p_dz, mask=(i_c * BTL + tl.arange(0, BTL)) < T) + + for _ in range(0, i_c * BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + # [BQ, BD] + b_dq += tl.dot((b_ds * (1 + b_s)).to(b_v.dtype), b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + + b_dq *= scale + o_q = tl.arange(0, BTL) + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_c * BTL), (BV, BTS), (0, 1)) + # Q block and K block have overlap. masks required + for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + b_s = tl.where(m_s, b_s, 0) + # [BTL, BK] + b_dq += tl.dot((b_ds + b_ds * b_s).to(b_k.dtype), b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + o_k += BTS + p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit +def _parallel_based_bwd_dkv( + i_bh, + i_c, + i_k, + i_v, + q, + k, + v, + do, + dz, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, +): + # compute dk dv + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) + b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) + b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros([BTL, BV], dtype=tl.float32) + + for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + b_q = tl.load(p_q, boundary_check=(0, 1)) # [BK, BTS] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) # [BV, BTS] + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale # [BTL, BTS] + b_s2 = 1 + b_s + 0.5 * b_s * b_s + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale + if i_v == 0: + b_ds += b_dz[None, :] * scale + else: + b_ds = b_ds + b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + + tl.debug_barrier() + o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) + for i in range(i_c*BTL, (i_c+1)*BTL, BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + # [BK, BQ] + m_s = o_k[:, None] <= o_q[None, :] + b_s = tl.dot(b_k, b_q, allow_tf32=False) * scale + b_s2 = 1 + b_s + 0.5 * b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_s2 = tl.where(m_s, b_s2, 0) + + b_ds = tl.dot(b_v, b_do, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[None, :] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + # [BK, BD] + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + o_q += BTS + + p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit(do_not_specialize=['T']) +def parallel_based_bwd_kernel( + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % NV + _parallel_based_bwd_dq( + i_bh, i_c, i_k, i_v, + q, k, v, do, dz, dq, + scale, T, B, H, BTL, BTS, BK, BV, K, V, + ) + tl.debug_barrier() + _parallel_based_bwd_dkv( + i_bh, i_c, i_k, i_v, + q, k, v, do, dz, dk, dv, + scale, T, B, H, BTL, BTS, BK, BV, K, V, + ) + + +class ParallelBasedFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale): + BTL, BTS = 128, 32 + assert BTL % BTS == 0 + # assert q.shape[-1] % 16 == 0 + BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not." + + o = torch.empty(NK, B, H, T, V, device=q.device) + z = torch.empty(NK, B, H, T, device=q.device) + parallel_based_fwd_kernel[grid]( + q, k, v, o, z, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages, + ) + ctx.save_for_backward(q, k, v) + ctx.scale = scale + return o.sum(0).to(q.dtype), z.sum(0).to(q.dtype) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dz): + q, k, v = ctx.saved_tensors + scale = ctx.scale + BTL, BTS = 64, 32 + assert BTL % BTS == 0 + BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not" + + dq = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dk = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dv = torch.empty(NK, B, H, T, V, dtype=q.dtype, device=q.device) + + parallel_based_bwd_kernel[grid]( + q, k, v, do, dz, dq, dk, dv, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages, + ) + + return dq.sum(0).to(q.dtype), dk.sum(0).to(k.dtype), dv.sum(0).to(v.dtype), None + + +triton_parallel_based = ParallelBasedFunction.apply + + +def parallel_based( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + use_norm: bool = True, + head_first: bool = False, +): + assert q.shape[-1] <= 128, "only support feature dim up to 128" + if scale is None: + scale = q.shape[-1] ** -0.5 + if not head_first: + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + o, z = triton_parallel_based(q, k, v, scale) + if use_norm: + o = o / (z[..., None] + 1e-6) + if not head_first: + o = o.transpose(1, 2) + return o.to(q.dtype) diff --git a/fla/ops/comba/__init__.py b/fla/ops/comba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba3ed4441292754caf416e0200deb20a14e89b51 --- /dev/null +++ b/fla/ops/comba/__init__.py @@ -0,0 +1,10 @@ +from .chunk import chunk_comba +from .fused_recurrent import fused_recurrent_comba +from .naive import naive_chunk_comba, naive_recurrent_comba + +__all__ = [ + "chunk_comba", + "fused_recurrent_comba", + "naive_chunk_comba", + "naive_recurrent_comba", +] diff --git a/fla/ops/comba/chunk.py b/fla/ops/comba/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..fb936db925a08560af0e3087250053eba2e1773d --- /dev/null +++ b/fla/ops/comba/chunk.py @@ -0,0 +1,380 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.comba.utils import chunk_comba_cumsum_scalar_bwd, chunk_comba_cumsum_scalar_fwd +from fla.ops.comba.wy_fast import chunk_scaled_dot_comba_pkt_fwd, prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices, solve_tril +from fla.ops.utils.constant import RCP_LN2 +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_comba_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + g0, g = chunk_comba_cumsum_scalar_fwd( + g, + chunk_size=64, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=RCP_LN2, + ) + # obtain WY representation. u is actually the new v. + A = chunk_scaled_dot_comba_pkt_fwd( + k=k, + p=p, + beta=beta, + g0=g0, + g=g, + cu_seqlens=cu_seqlens, + output_dtype=torch.float32, + chunk_indices=chunk_indices, + use_exp2=True, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + w, u = recompute_w_u_fwd( + k=p, + v=v, + beta=beta, + A=A, + g_cumsum=g0, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + return g0, g, o, A, final_state + + +def chunk_comba_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g0: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + w, u = recompute_w_u_fwd( + k=p, + v=v, + beta=beta, + A=A, + g_cumsum=g0, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dk2, dv, dp, db, dg0, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + p=p, + beta=beta, + g0=g0, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dk.add_(dk2) + dg.add_(dg2) + assert dg.dtype == torch.float32, "dg should be fp32" + dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + # dg0 = d(g_cumsum - g) + dg += chunk_comba_cumsum_scalar_bwd(dg0, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + return dq, dk, dv, dp, db, dg, dh0 + + +class ChunkCombaFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + ): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + p, p_rstd = l2norm_fwd(p) + else: + q_rstd, k_rstd, p_rstd = None, None, None + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g0, g, o, A, final_state = chunk_comba_fwd( + q=q, + k=k, + v=v, + p=p, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, p, p_rstd, v, g0, g, beta, A, initial_state, cu_seqlens, + chunk_indices) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + q, q_rstd, k, k_rstd, p, p_rstd, v, g0, g, beta, A, initial_state, cu_seqlens, chunk_indices = ( + ctx.saved_tensors + ) + dq, dk, dv, dp, db, dg, dh0 = chunk_comba_bwd( + q=q, + k=k, + v=v, + p=p, + g0=g0, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + dp = l2norm_bwd(p, p_rstd, dp) + return dq.to(q), dk.to(k), dv.to(v), dp.to(p), dg.to(g), db.to(beta), None, dh0, None, None, None, None + + +@torch.compiler.disable +def chunk_comba( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + p (torch.Tensor): + auxiliary keys of shape `[B, T, H, K]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.comba import chunk_comba + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> b = torch.rand(H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> p = k * b[:, None] + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_comba( + q, k, v, p, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = chunk_comba( + q, k, v, p, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if p is None: + p = k + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkCombaFunction.apply( + q, + k, + v, + p, + g, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/comba/fused_recurrent.py b/fla/ops/comba/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..13c150090966b84406bec5ef88b9c3c32b37e1bd --- /dev/null +++ b/fla/ops/comba/fused_recurrent.py @@ -0,0 +1,335 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, exp2 +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_comba_fwd_kernel( + q, + k, + p, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + STORE_FINAL_STATE: tl.constexpr, # whether to store final state + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_EXP2: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + p_p = p + (bos * H + i_h) * K + o_k + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + p_g = g + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_p = tl.load(p_p, mask=mask_k, other=0).to(tl.float32) + b_g = tl.load(p_g).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_p = b_p / tl.sqrt(tl.sum(b_p * b_p) + 1e-6) + b_q = b_q * scale + # [BV] + b_v -= tl.sum(b_h * b_p[:, None], 0) + # [BK, BV] + if USE_EXP2: + b_h *= exp2(b_g) + else: + b_h *= exp(b_g) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_o += HV*V + p_v += HV*V + p_p += H*K + p_g += HV + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_comba_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + o = q.new_empty(NK, *v.shape) + if output_final_state: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NK, NV, N * HV) + fused_recurrent_comba_fwd_kernel[grid]( + q=q, + k=k, + p=p, + v=v, + g=g, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_EXP2=use_exp2, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o, final_state + + +class FusedRecurrentCombaFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + p: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, final_state = fused_recurrent_comba_fwd( + q=q, + k=k, + p=p, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps.", + ) + + +def fused_recurrent_comba( + q: torch.Tensor, + k: torch.Tensor, + p: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + p (torch.Tensor): + auxiliary keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use qk l2norm within the kernel for saving GPU memory. + Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.comba import fused_recurrent_comba + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> b = torch.rand(H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> p = k * b[:, None] + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_recurrent_comba( + q, k, v, p, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, p, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, p, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_comba( + q, k, p, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if beta is None: + beta = torch.ones_like(q[..., 0]) + if p is None: + p = k + o, final_state = FusedRecurrentCombaFunction.apply( + q, + k, + p, + v, + g, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/comba/naive.py b/fla/ops/comba/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..173339e3cac513d12bb19720c75e50855ba4eeee --- /dev/null +++ b/fla/ops/comba/naive.py @@ -0,0 +1,166 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_recurrent_comba( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of recurrent COMBA. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + p: [B, T, H, K] + beta: [B, T, H] + g: [B, T, H] + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g]) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + b_p = p[:, :, i] + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_p[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def naive_chunk_comba( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of chunk COMBA. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + p: [B, T, H, K] + g: [B, T, H] + beta: [B, T, H] + chunk_size: int + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + p = F.pad(p, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + p_beta = p * beta[..., None] + assert l % chunk_size == 0 + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, p_beta, decay, g = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, p_beta, decay.unsqueeze(-1), g.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) # [B, H, n, c] + decay_0 = decay - g.squeeze(-1) # [B, H, n, c] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + L_mask_0 = ((decay_0.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + + # [B, H, n, c, d] @ [B, H, n, d, c] -> [B, H, n, c, c] + attn = -((p_beta @ k.transpose(-1, -2)) * L_mask_0).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + + # for U + k_cumsum = attn @ v + # for W + k_cumdecay = attn @ (p_beta * decay_0[..., None].exp()) + v = k_cumsum + + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S = initial_state.to(torch.float32) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = k_cumdecay[:, :, i] @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S diff --git a/fla/ops/comba/utils.py b/fla/ops/comba/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..30aa090c5e619d9f0aa7020524c8c23cb984ca3b --- /dev/null +++ b/fla/ops/comba/utils.py @@ -0,0 +1,164 @@ + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['B', 'H', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_comba_cumsum_scalar_fwd_kernel( + g, + g0, + g1, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g1 = tl.make_block_ptr(g1 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if HAS_SCALE: + b_g = b_g * scale + b_g1 = tl.cumsum(b_g, axis=0) + b_g0 = b_g1 - b_g + tl.store(p_g0, b_g0.to(p_g0.dtype.element_ty), boundary_check=(0,)) + tl.store(p_g1, b_g1.to(p_g1.dtype.element_ty), boundary_check=(0,)) + + +def chunk_comba_cumsum_scalar_fwd( + g: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, + scale: float | None = None, +) -> torch.Tensor: + B, T, H = g.shape + assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + g0, g1 = torch.empty_like(g, dtype=output_dtype or g.dtype), torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = (NT, B * H) + chunk_comba_cumsum_scalar_fwd_kernel[grid]( + g, + g0, + g1, + scale, + cu_seqlens, + chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + ) + return g0, g1 + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['B', 'H', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_comba_cumsum_scalar_bwd_kernel( + dg0, + dgr, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + """ + b_dg: 1,2,3,4 + b_dg0: 0,1,2,3 + b_temp: 0,1,3,6 + b_dz: 6 + b_dgr: 6,5,3,0 + """ + b_dg0 = tl.load(p_dg0, boundary_check=(0,)).to(tl.float32) + b_temp = tl.cumsum(b_dg0, axis=0) + b_dz = tl.sum(b_dg0, axis=0) + b_dgr = -b_temp + b_dz[None] + tl.store(p_dgr, b_dgr.to(p_dgr.dtype.element_ty), boundary_check=(0,)) + + +def chunk_comba_cumsum_scalar_bwd( + dg0: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H = dg0.shape + assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + dg = torch.empty_like(dg0, dtype=output_dtype or dg0.dtype) + grid = (NT, B * H) + chunk_comba_cumsum_scalar_bwd_kernel[grid]( + dg0, + dg, + cu_seqlens, + chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + ) + return dg diff --git a/fla/ops/comba/wy_fast.py b/fla/ops/comba/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..b6cd1fade608eb3cf139654eefd840ad26108851 --- /dev/null +++ b/fla/ops/comba/wy_fast.py @@ -0,0 +1,450 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, exp2 +from fla.utils import autotune_cache_kwargs, check_shared_mem + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_G': lambda args: args['g'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'BT', 'IS_VARLEN', 'USE_G'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_comba_pkt_fwd_kernel( + k, + p, + beta, + g0, + g, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_p = tl.load(p_p, boundary_check=(0, 1)) + b_pb = b_p * b_beta[:, None] + b_A += tl.dot(b_pb.to(b_k.dtype), tl.trans(b_k)) + + if USE_G: + p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g0 = tl.load(p_g0, boundary_check=(0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + if USE_EXP2: + b_A = b_A * exp2(b_g0[:, None] - b_g[None, :]) + else: + b_A = b_A * exp(b_g0[:, None] - b_g[None, :]) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_comba_pkt_fwd( + k: torch.Tensor, + p: torch.Tensor, + beta: torch.Tensor, + g0: torch.Tensor | None = None, + g: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> torch.Tensor: + r""" + Compute beta \mathcal{A}(i-1/j) * P * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + p (torch.Tensor): + The auxiliary key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + g0 (torch.Tensor): + The cumulative sum minus the original one of the gate tensor of shape `[B, T, H]`. + Default: None + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. + Default: None + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_comba_pkt_fwd_kernel[(NT, B * H)]( + k=k, + p=p, + beta=beta, + g0=g0, + g=g, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + USE_EXP2=use_exp2, + ) + return A + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + p, + beta, + g0, + g, + A, + dw, + du, + dk, + dv, + dp, + dbeta, + dg0, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_EXP2: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g0 = tl.make_block_ptr(g0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_g0 = tl.load(p_g0, boundary_check=(0,)) + if USE_EXP2: + b_g0_exp = exp2(b_g0) + else: + b_g0_exp = tl.exp(b_g0) + b_g = tl.load(p_g, boundary_check=(0,)) + + b_dbeta = tl.zeros([BT], dtype=tl.float32) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_dg0 = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_p = tl.load(p_p, boundary_check=(0, 1)) + b_p_beta_g0 = (b_p * b_beta[:, None] * b_g0_exp[:, None]).to(b_p.dtype) + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dA += tl.dot(b_dw, tl.trans(b_p_beta_g0)) + b_dp_beta_g0 = tl.dot(b_A, b_dw) + b_dp = b_dp_beta_g0 * b_beta[:, None] * b_g0_exp[:, None] + b_dbeta += tl.sum(b_dp_beta_g0 * b_p * b_g0_exp[:, None], 1) + b_dg0 += tl.sum(b_dp * b_p, 1) + tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_v_beta)) + b_dv_beta = tl.dot(b_A, b_du) + b_dv = b_dv_beta * b_beta[:, None] + b_dbeta += tl.sum(b_dv_beta * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + if USE_EXP2: + b_dA = tl.where(m_A, -b_dA * exp2(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) + else: + b_dA = tl.where(m_A, -b_dA * exp(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) + b_dA = b_dA.to(k.dtype.element_ty) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_p = tl.load(p_p, boundary_check=(0, 1)) + b_dp = tl.load(p_dp, boundary_check=(0, 1)) + b_p_beta = (b_p * b_beta[:, None]).to(b_p.dtype) + b_A += tl.dot(b_p_beta, tl.trans(b_k)) + b_dp_beta = tl.dot(b_dA, b_k) + b_dbeta += tl.sum(b_dp_beta * b_p, 1) + b_dk = tl.dot(tl.trans(b_dA), b_p_beta) + b_dp += b_dp_beta * b_beta[:, None] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) + + b_dA_A = b_dA * b_A + b_dg0 += tl.sum(b_dA_A, axis=1) + b_dg = - tl.sum(b_dA_A, axis=0) + p_dg = tl.make_block_ptr(dg + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dg0 = tl.make_block_ptr(dg0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dbeta = tl.make_block_ptr(dbeta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dg0, b_dg0.to(p_dg0.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_EXP2: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + if USE_EXP2: + b_g = exp2(tl.load(p_g, boundary_check=(0,))) + else: + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A, b_kb) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g_cumsum: torch.Tensor, + A: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = 64 + BV = 64 + + u = torch.empty_like(v) + w = torch.empty_like(k) + recompute_w_u_fwd_kernel[(NT, B*H)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g_cumsum, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + ) + return w, u + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g0: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dp = torch.empty_like(p) + dbeta = torch.empty_like(beta) + dg0 = torch.empty_like(g0) + dg = torch.empty_like(g) + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + k=k, + v=v, + p=p, + beta=beta, + g0=g0, + g=g, + A=A, + dw=dw, + du=du, + dk=dk, + dv=dv, + dp=dp, + dbeta=dbeta, + dg0=dg0, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + ) + return dk, dv, dp, dbeta, dg0, dg diff --git a/fla/ops/common/__init__.py b/fla/ops/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/fla/ops/common/backends/__init__.py b/fla/ops/common/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72b42739a76b5800315158bf6b2633c6663f73fc --- /dev/null +++ b/fla/ops/common/backends/__init__.py @@ -0,0 +1,12 @@ +"""Common backends for shared operations like chunk_gated_delta_rule_fwd_h.""" + +from fla.ops.backends import BackendRegistry, dispatch +from fla.ops.common.backends.intracard import IntraCardCPBackend + +common_registry = BackendRegistry("common") + + +common_registry.register(IntraCardCPBackend()) + + +__all__ = ['common_registry', 'dispatch'] diff --git a/fla/ops/common/backends/intracard.py b/fla/ops/common/backends/intracard.py new file mode 100644 index 0000000000000000000000000000000000000000..195b2d818a51a9a76feb1479df7c8ee3b0ba67e7 --- /dev/null +++ b/fla/ops/common/backends/intracard.py @@ -0,0 +1,94 @@ +"""Intra-card CP backend for shared delta rule operations. + +Accelerates prefill by splitting long sequences into sub-sequences +and processing them in parallel across SMs. + +Only active under torch.inference_mode() with varlen (cu_seqlens != None). +""" + +from __future__ import annotations + +import os + +import torch + +from fla.ops.backends import BaseBackend + +# Maximum number of sub-sequences per original sequence +# Limits merge chain depth to control precision loss +MAX_SUBSEQS = int(os.environ.get('FLA_INTRACARD_MAX_SPLITS', 32)) + + +class IntraCardCPBackend(BaseBackend): + """Intra-card context parallel backend for chunk_gated_delta_rule_fwd_h.""" + + backend_type = "intracard_cp" + package_name = None # No external package needed + env_var = "FLA_INTRACARD_CP" + default_enable = False + + @classmethod + def is_available(cls) -> bool: + return True + + def chunk_gated_delta_rule_fwd_h_verifier( + self, + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, + ) -> tuple[bool, str | None]: + """Check if intracard CP should handle this call.""" + # Only in inference mode + if not torch.is_inference_mode_enabled(): + return False, "Not in inference mode" + + # Only for varlen + if cu_seqlens is None: + return False, "cu_seqlens is None" + + return True, None + + def chunk_gated_delta_rule_fwd_h( + self, + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Intra-card CP implementation of chunk_gated_delta_rule_fwd_h.""" + from fla.ops.common.intracard_cp import intracard_fwd_h + + return intracard_fwd_h( + k=k, w=w, u=u, g=g, gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + max_splits=MAX_SUBSEQS, + transpose_state_layout=transpose_state_layout, + ) diff --git a/fla/ops/common/chunk_delta_h.py b/fla/ops/common/chunk_delta_h.py new file mode 100644 index 0000000000000000000000000000000000000000..727a4b2dce90ee6ffb8d2db7d26057b209470c6c --- /dev/null +++ b/fla/ops/common/chunk_delta_h.py @@ -0,0 +1,769 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp, exp2 +from fla.utils import IS_NVIDIA_HOPPER, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem('ampere') else [2, 1]) + for BV in ([32, 64] if check_shared_mem('ada') else [32]) + ], + key=['H', 'K', 'V', 'BT', 'USE_EXP2', 'TRANSPOSE_STATE'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + USE_EXP2: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if TRANSPOSE_STATE: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * H + i_h).to(tl.int64) * K*V + v += (bos * H + i_h).to(tl.int64) * V + k += (bos * H + i_h).to(tl.int64) * K + w += (bos * H + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * H + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K*V + if STORE_FINAL_STATE: + ht = ht + i_nh * K*V + + # load initial state + if USE_INITIAL_STATE: + if TRANSPOSE_STATE: + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + if TRANSPOSE_STATE: + p_h0_2 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + if TRANSPOSE_STATE: + p_h0_3 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + if TRANSPOSE_STATE: + p_h0_4 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + if TRANSPOSE_STATE: + p_h1 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_h1 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + if TRANSPOSE_STATE: + p_h2 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_h2 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + if TRANSPOSE_STATE: + p_h3 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_h3 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + if TRANSPOSE_STATE: + p_h4 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_h4 = tl.make_block_ptr(h + i_t_int64 * H*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + else: + b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + else: + b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + else: + b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + if SAVE_NEW_VALUE: + p_v = tl.make_block_ptr(v_new, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + (bos * H + last_idx * H + i_h).to(tl.int64)).to(tl.float32) + p_g = tl.make_block_ptr(g + (bos * H + i_h).to(tl.int64), (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + if USE_EXP2: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp(b_gk_last1)[None, :] + else: + if USE_EXP2: + b_h1 *= exp2(b_gk_last1)[:, None] + else: + b_h1 *= exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + if USE_EXP2: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp(b_gk_last2)[None, :] + else: + if USE_EXP2: + b_h2 *= exp2(b_gk_last2)[:, None] + else: + b_h2 *= exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + if USE_EXP2: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp(b_gk_last3)[None, :] + else: + if USE_EXP2: + b_h3 *= exp2(b_gk_last3)[:, None] + else: + b_h3 *= exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + if USE_EXP2: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp(b_gk_last4)[None, :] + else: + if USE_EXP2: + b_h4 *= exp2(b_gk_last4)[:, None] + else: + b_h4 *= exp(b_gk_last4)[:, None] + + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 += tl.dot(b_k, b_v) + + if STORE_FINAL_STATE: + if TRANSPOSE_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + if TRANSPOSE_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + if TRANSPOSE_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + if TRANSPOSE_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + USE_EXP2: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if TRANSPOSE_STATE: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h).to(tl.int64) * K + k += (bos * H + i_h).to(tl.int64) * K + w += (bos * H + i_h).to(tl.int64) * K + do += (bos * H + i_h).to(tl.int64) * V + dv += (bos * H + i_h).to(tl.int64) * V + dv2 += (bos * H + i_h).to(tl.int64) * V + dh += (boh * H + i_h).to(tl.int64) * K*V + if USE_GK: + gk += (bos * H + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K*V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K*V + + if USE_FINAL_STATE_GRADIENT: + if TRANSPOSE_STATE: + p_dht1 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + if K > 64: + if TRANSPOSE_STATE: + p_dht2 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + if K > 128: + if TRANSPOSE_STATE: + p_dht3 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + if K > 192: + if TRANSPOSE_STATE: + p_dht4 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + if TRANSPOSE_STATE: + p_dh1 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_dh1 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + if TRANSPOSE_STATE: + p_dh2 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_dh2 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + if TRANSPOSE_STATE: + p_dh3 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_dh3 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + if TRANSPOSE_STATE: + p_dh4 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_dh4 = tl.make_block_ptr(dh + i_t_int64*H*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + else: + bg_last_exp = exp(bg_last) + b_g_exp = exp(b_g) + + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_do = tl.load(p_do, boundary_check=(0, 1)) + + # Update dv + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * H*K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk + last_idx * H*K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + b_dv += tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype)) + else: + b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + + if K > 128: + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load(gk + last_idx * H*K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + b_dv += tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype)) + else: + b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + + if K > 192: + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load(gk + last_idx * H*K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + if TRANSPOSE_STATE: + b_dv += tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype)) + else: + b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + if USE_EXP2: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + else: + b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, boundary_check=(0, 1)) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + # Update dh + p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if TRANSPOSE_STATE: + if USE_EXP2: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp(b_gk_last1)[None, :] + else: + if USE_EXP2: + b_dh1 *= exp2(b_gk_last1[:, None]) + else: + b_dh1 *= exp(b_gk_last1[:, None]) + if TRANSPOSE_STATE: + b_dh1 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 64: + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if TRANSPOSE_STATE: + if USE_EXP2: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp(b_gk_last2)[None, :] + else: + if USE_EXP2: + b_dh2 *= exp2(b_gk_last2[:, None]) + else: + b_dh2 *= exp(b_gk_last2[:, None]) + if TRANSPOSE_STATE: + b_dh2 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 128: + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if TRANSPOSE_STATE: + if USE_EXP2: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp(b_gk_last3)[None, :] + else: + if USE_EXP2: + b_dh3 *= exp2(b_gk_last3[:, None]) + else: + b_dh3 *= exp(b_gk_last3[:, None]) + if TRANSPOSE_STATE: + b_dh3 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 192: + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if TRANSPOSE_STATE: + if USE_EXP2: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp(b_gk_last4)[None, :] + else: + if USE_EXP2: + b_dh4 *= exp2(b_gk_last4[:, None]) + else: + b_dh4 *= exp(b_gk_last4[:, None]) + if TRANSPOSE_STATE: + b_dh4 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + if USE_INITIAL_STATE: + if TRANSPOSE_STATE: + p_dh0 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + else: + p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + if TRANSPOSE_STATE: + p_dh1 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + else: + p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + if TRANSPOSE_STATE: + p_dh2 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + else: + p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + if TRANSPOSE_STATE: + p_dh3 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + else: + p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + + +@dispatch('common') +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V = *k.shape, u.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if transpose_state_layout: + h = k.new_empty(B, NT, H, V, K) + final_state = k.new_zeros(N, H, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_zeros(N, H, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + ) + return h, v_new, final_state + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *q.shape, do.shape[-1] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = 64 + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + if transpose_state_layout: + dh = q.new_empty(B, NT, H, V, K) + else: + dh = q.new_empty(B, NT, H, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BV=64 if check_shared_mem('ada', q.device.index) else 32, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + # H200 autotune spends a long time walking this backward state kernel + # across both ranks. Use a deterministic fixed launch. + num_warps=4, + num_stages=2, + ) + return dh, dh0, dv2 diff --git a/fla/ops/common/chunk_h.py b/fla/ops/common/chunk_h.py new file mode 100644 index 0000000000000000000000000000000000000000..5c29d999a81eb62e0a55ae40cd0d46bf7108535a --- /dev/null +++ b/fla/ops/common/chunk_h.py @@ -0,0 +1,384 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, check_shared_mem + +BKV_LIST = [32, 64] if check_shared_mem() else [16, 32] + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_h( + k, + v, + h, + g, + g_gamma, + gk, + gv, + h0, + ht, + cu_seqlens, + split_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) + boh = tl.load(split_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) + boh = i_n * NS + NTS = BS // BT + + if USE_G_GAMMA: + # decay rate given the head index + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + i_s = i_t // NTS + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V + p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + if i_t % NTS == 0: + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + last_idx = min((i_t + 1) * BT, T) - 1 + + # scalar decay + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_h *= exp(b_g_last) + b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) + + if USE_G_GAMMA: + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_h *= exp(b_g_last) + b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) + + # vector decay, h = Diag(gk) @ h + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_h *= exp(b_gk_last)[:, None] + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) + + # vector decay, h = h @ Diag(gv) + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_h *= exp(b_gv_last)[None, :] + + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) + + b_h += tl.dot(b_k, b_v) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dh( + q, + g, + g_gamma, + gk, + gv, + do, + dh, + dht, + dh0, + cu_seqlens, + split_offsets, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // NG + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + NS = tl.cdiv(T, BS) + boh = tl.load(split_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + NS = tl.cdiv(T, BS) + boh = i_n * NS + + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT - 1, -1, -1): + i_s = i_t // (BS // BT) + o_dh = ((boh + i_s) * H + i_h).to(tl.int64) * K*V + p_dh = tl.make_block_ptr(dh + o_dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + if i_t % (BS // BT) == 0: + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + # [BK, BT] + p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + if USE_G: + p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) + b_dh *= exp(b_g_last) + + if USE_G_GAMMA: + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) + b_dh *= exp(b_g_last) + + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_q = (b_q * exp(b_gk)).to(b_q.dtype) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_dh *= exp(b_gk_last)[:, None] + + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_do = (b_do * exp(b_gv)) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_dh *= exp(b_gv_last)[None, :] + + b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + + if STORE_INITIAL_STATE_GRADIENT: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + split_size: int | None = None, + states_in_fp32: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + BS = BT if split_size is None else split_size + assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NS, split_offsets = B, triton.cdiv(T, BS), None + else: + split_offsets = prepare_chunk_offsets(cu_seqlens, BS) + N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() + + h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + chunk_fwd_kernel_h[grid]( + k=k, + v=v, + h=h, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + h0=h0, + ht=ht, + cu_seqlens=cu_seqlens, + split_offsets=split_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + USE_G=g is not None, + USE_G_GAMMA=g_gamma is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + return h, ht + + +def chunk_bwd_dh( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor, + scale: float, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + split_size: int | None = None, + states_in_fp32: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + BT = chunk_size + BS = BT if split_size is None else split_size + assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" + # N: the actual number of sequences in the batch with either equal or variable lengths + # NG: number of groups in GQA + if cu_seqlens is None: + N, NS, split_offsets = B, triton.cdiv(T, BS), None + else: + split_offsets = prepare_chunk_offsets(cu_seqlens, BS) + N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() + NG = HQ // H + + dh = k.new_empty(B, NS, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None + + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + chunk_bwd_kernel_dh[grid]( + q=q, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + do=do, + dh=dh, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + split_offsets=split_offsets, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + NG=NG, + USE_G=g is not None, + USE_G_GAMMA=g_gamma is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + return dh, dh0 diff --git a/fla/ops/common/chunk_h_parallel.py b/fla/ops/common/chunk_h_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..77f86ad02348298cdb80f4e22b042041cc2deb05 --- /dev/null +++ b/fla/ops/common/chunk_h_parallel.py @@ -0,0 +1,558 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +Fully parallelized state passing. +""" + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for BV in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_h_parallel( + k, + v, + h, + g, + gk, + gv, + h0, + ht, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + NV = tl.cdiv(V, BV) + # i_b: batch index + # i_h: head index + # i_n: sequence index + # i_t: chunk index within current sequence + # i_tg: (global) chunk index across all sequences + i_k, i_v = i_kv // NV, i_kv % NV + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + bos, eos = i_b * T, i_b * T + T + NT = tl.cdiv(T, BT) + i_n, i_tg = i_b, i_b * NT + i_t + i_nh = i_n * H + i_h + + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + if i_t == 0: + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + last_idx = min(i_t * BT + BT, T) - 1 + # scalar decay + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) + + # vector decay, h = Diag(gk) @ h + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) + + # vector decay, h = h @ Diag(gv) + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) + + b_h = tl.dot(b_k, b_v) + if i_t < NT - 1: + p_h = tl.make_block_ptr(h + ((i_tg + 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + elif STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for BV in [32, 64, 128] + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_h_reduction( + h, + g, + gk, + gv, + kvt, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + for i_t in range(NT): + p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + if i_t > 0: + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min(i_t * BT + BT, T) - 1 + # scalar decay + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + b_h *= exp(b_g_last) + + # vector decay, h = Diag(gk) @ h + if USE_GK: + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_h *= exp(b_gk_last)[:, None] + + # vector decay, h = h @ Diag(gv) + if USE_GV: + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_h *= exp(b_gv_last)[None, :] + + if STORE_FINAL_STATE: + p_kvt = tl.make_block_ptr(kvt + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_kvt, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for BV in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dh_parallel( + q, + g, + gk, + gv, + do, + dh, + dht, + dh0, + cu_seqlens, + chunk_indices, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + NV = tl.cdiv(V, BV) + i_k, i_v = i_kv // NV, i_kv % NV + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + bos, eos = i_b * T, i_b * T + T + NT = tl.cdiv(T, BT) + i_n, i_tg = i_b, i_b * NT + i_t + i_nh = i_n * HQ + i_hq + + p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + if i_t == NT - 1: + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh = tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + else: + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + + # [BK, BT] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + if USE_G: + p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) + + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_q = (b_q * exp(b_gk)).to(b_q.dtype) + + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_do = (b_do * exp(b_gv)).to(b_do.dtype) + + b_dh = tl.dot(b_q, b_do) + if i_t > 0: + p_dh = tl.make_block_ptr(dh + ((i_tg - 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + elif STORE_INITIAL_STATE_GRADIENT: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for BV in [32, 64, 128] + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dh_reduction( + g, + gk, + gv, + dh, + doq0, + dh0, + cu_seqlens, + chunk_offsets, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // NG + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + for i_t in range(NT - 1, -1, -1): + p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) + if i_t < NT - 1: + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min(i_t * BT + BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + b_dh *= exp(b_g_last) + + if USE_GK: + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_dh *= exp(b_gk_last)[:, None] + + if USE_GV: + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_dh *= exp(b_gv_last)[None, :] + + if STORE_INITIAL_STATE_GRADIENT: + p_doq0 = tl.make_block_ptr(doq0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_doq0, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + gk: torch.Tensor, + gv: torch.Tensor, + h0: torch.Tensor, + output_final_state: bool, + states_in_fp32: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + h = k.new_empty(B, NT, H, K, V, dtype=torch.float) + ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * H) + chunk_fwd_kernel_h_parallel[grid]( + k=k, + v=v, + h=h, + g=g, + gk=gk, + gv=gv, + h0=h0, + ht=ht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + kvt, ht = ht, (torch.empty_like(ht) if output_final_state else None) + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + chunk_fwd_kernel_h_reduction[grid]( + h=h, + g=g, + gk=gk, + gv=gv, + kvt=kvt, + ht=ht, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + h = h.to(k.dtype) if not states_in_fp32 else h + return h, ht + + +def chunk_bwd_dh( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + gk: torch.Tensor, + gv: torch.Tensor, + do: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor, + scale: float, + states_in_fp32: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # N: the actual number of sequences in the batch with either equal or variable lengths + # NG: number of groups in GQA + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + NG = HQ // H + + dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None + + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * HQ) + chunk_bwd_kernel_dh_parallel[grid]( + q=q, + g=g, + gk=gk, + gv=gv, + do=do, + dh=dh, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + + doq0, dh0 = dh0, (torch.empty_like(dh0) if dh0 is not None else None) + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) + chunk_bwd_kernel_dh_reduction[grid]( + g=g, + gk=gk, + gv=gv, + dh=dh, + doq0=doq0, + dh0=dh0, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + dh = dh.to(q.dtype) if not states_in_fp32 else dh + return dh, dh0 diff --git a/fla/ops/common/chunk_h_split.py b/fla/ops/common/chunk_h_split.py new file mode 100644 index 0000000000000000000000000000000000000000..8f601eb89c9f794b0d69fa07761aa3ffca31d649 --- /dev/null +++ b/fla/ops/common/chunk_h_split.py @@ -0,0 +1,599 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_h_split( + k, + v, + g, + gk, + gv, + hs, + hr, + h0, + ht, + cu_seqlens, + split_indices, + T, + S: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # handle one split at a time + # i_h: head index + # i_n: sequence index + # i_s: local split index inside a sequence + i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_ss, i_h = i_sh // H, i_sh % H + if IS_VARLEN: + i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + else: + NS = tl.cdiv(T, S) + i_n, i_s = i_ss // NS, i_ss % NS + bos, eos = i_n * T, i_n * T + T + i_nh = i_n * H + i_h + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # for the first split, we directly store the state as the final result + if i_s == 0: + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_hr = tl.make_block_ptr(hr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) + for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + + # scalar decay + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h + b_h *= exp(b_g_last) + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) + + # vector decay, h = Diag(gk) @ h + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_h *= exp(b_gk_last)[:, None] + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) + + # vector decay, h = h @ Diag(gv) + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_h *= exp(b_gv_last)[None, :] + + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) + + b_h += tl.dot(b_k, b_v) + + # if there are more than one splits, we store the result to (unreduced) hs + # otherwise, we store the result to ht as the final state + if NS > 1: + p_hs = tl.make_block_ptr(hs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), boundary_check=(0, 1)) + elif STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_h_reduction( + g, + gk, + gv, + hs, + hr, + ht, + cu_seqlens, + split_offsets, + T, + S: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + boh = tl.load(split_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NS = tl.cdiv(T, S) + boh = i_n * NS + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # skip the first split + for i_s in range(1, NS): + p_hs = tl.make_block_ptr(hs + ((boh + i_s-1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_hr = tl.make_block_ptr(hr + ((boh + i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) + + for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): + last_idx = min(i_t * BT + BT, T) - 1 + # scalar decay + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + b_h *= exp(b_g_last) + + # vector decay, h = Diag(gk) @ h + if USE_GK: + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_h *= exp(b_gk_last)[:, None] + + # vector decay, h = h @ Diag(gv) + if USE_GV: + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_h *= exp(b_gv_last)[None, :] + + if NS > 1: + if STORE_FINAL_STATE: + p_hs = tl.make_block_ptr(hs + ((boh + NS-1) * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dh_split( + q, + g, + gk, + gv, + do, + dht, + dhs, + dhr, + dh0, + cu_seqlens, + split_indices, + scale, + T, + S: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # handle one split at a time + # i_h: head index + # i_n: sequence index + # i_s: local split index inside a sequence + i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_ss, i_hq = i_sh // HQ, i_sh % HQ + if IS_VARLEN: + i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + else: + NS = tl.cdiv(T, S) + i_n, i_s = i_ss // NS, i_ss % NS + bos, eos = i_n * T, i_n * T + T + i_nh = i_n * HQ + i_hq + i_h = i_hq // NG + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if i_s == NS - 1: + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + p_dhr = tl.make_block_ptr(dhr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) + + for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): + p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + last_idx = min(i_t * BT + BT, T) - 1 + if USE_G: + p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) + b_dh *= exp(b_g_last) + + if USE_GK: + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_q = (b_q * exp(b_gk)).to(b_q.dtype) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_dh *= exp(b_gk_last)[:, None] + + if USE_GV: + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_do = (b_do * exp(b_gv)).to(b_do.dtype) + + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_dh *= exp(b_gv_last)[None, :] + + b_dh += tl.dot(b_q, b_do) + + if NS > 1: + p_dhs = tl.make_block_ptr(dhs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dhs, b_dh.to(p_dhs.dtype.element_ty), boundary_check=(0, 1)) + elif STORE_INITIAL_STATE_GRADIENT: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dh_reduction( + g, + gk, + gv, + dhs, + dhr, + dh0, + cu_seqlens, + split_offsets, + T, + S: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // NG + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + boh = tl.load(split_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NS = tl.cdiv(T, S) + boh = i_n * NS + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + for i_s in range(NS - 2, -1, -1): + p_dhs = tl.make_block_ptr(dhs + ((boh+i_s+1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dhr = tl.make_block_ptr(dhr + ((boh+i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) + + for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): + last_idx = min(i_t * BT + BT, T) - 1 + # scalar decay + if USE_G: + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + b_dh *= exp(b_g_last) + + if USE_GK: + p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_dh *= exp(b_gk_last)[:, None] + + if USE_GV: + p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) + b_dh *= exp(b_gv_last)[None, :] + + if NS > 1: + if STORE_INITIAL_STATE_GRADIENT: + p_dhs = tl.make_block_ptr(dhs + (boh * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + gk: torch.Tensor, + gv: torch.Tensor, + h0: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + split_offsets: torch.LongTensor | None = None, + split_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + split_size: int = 256, + states_in_fp32: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + # B: batch size + # N: the actual number of sequences in the batch + # H: number of heads + # T: sequence length, can be variable across sequences + # S: split size, a multiple of chunk size + # BT: chunk size + S, BT = split_size, chunk_size + assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" + if cu_seqlens is None: + N = B + NS = N * triton.cdiv(T, S) + else: + N = len(cu_seqlens) - 1 + NS = split_offsets[-1] + + # unreduced kv states per split + hs = k.new_empty(NS, H, K, V, dtype=torch.float) + # reduced states per split + hr = k.new_empty(NS, H, K, V, dtype=torch.float if states_in_fp32 else k.dtype) + ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + # parallelized over splits + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * H) + chunk_fwd_kernel_h_split[grid]( + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + hs=hs, + hr=hr, + h0=h0, + ht=ht, + cu_seqlens=cu_seqlens, + split_indices=split_indices, + T=T, + S=S, + H=H, + K=K, + V=V, + BT=BT, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + chunk_fwd_kernel_h_reduction[grid]( + g=g, + gk=gk, + gv=gv, + hs=hs, + hr=hr, + ht=ht, + cu_seqlens=cu_seqlens, + split_offsets=split_offsets, + T=T, + S=S, + H=H, + K=K, + V=V, + BT=BT, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + return hr, ht + + +def chunk_bwd_dh( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + gk: torch.Tensor, + gv: torch.Tensor, + do: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + split_offsets: torch.Tensor | None = None, + split_indices: torch.Tensor | None = None, + chunk_size: int = 64, + split_size: int = 256, + states_in_fp32: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + # B: batch size + # N: the actual number of sequences in the batch + # H: number of heads + # T: sequence length, can be variable across sequences + # S: split size, a multiple of chunk size + # BT: chunk size + S, BT = max(chunk_size, min(split_size, triton.next_power_of_2(T))), chunk_size + assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" + if cu_seqlens is None: + N = B + NS = N * triton.cdiv(T, S) + else: + N = len(cu_seqlens) - 1 + NS = split_offsets[-1] + # number of groups in GQA + NG = HQ // H + + dhs = q.new_empty(NS, HQ, K, V, dtype=torch.float) + dhr = q.new_empty(NS, HQ, K, V, dtype=torch.float if states_in_fp32 else k.dtype) + dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None + + # parallelized over splits + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * HQ) + chunk_bwd_kernel_dh_split[grid]( + q=q, + g=g, + gk=gk, + gv=gv, + do=do, + dht=dht, + dhs=dhs, + dhr=dhr, + dh0=dh0, + cu_seqlens=cu_seqlens, + split_indices=split_indices, + scale=scale, + T=T, + S=S, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) + chunk_bwd_kernel_dh_reduction[grid]( + g=g, + gk=gk, + gv=gv, + dhs=dhs, + dhr=dhr, + dh0=dh0, + cu_seqlens=cu_seqlens, + split_offsets=split_offsets, + T=T, + S=S, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + USE_G=g is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + ) + return dhr, dh0 diff --git a/fla/ops/common/chunk_o.py b/fla/ops/common/chunk_o.py new file mode 100644 index 0000000000000000000000000000000000000000..43471321e00b26d4875d2b3d84a36984994c8184 --- /dev/null +++ b/fla/ops/common/chunk_o.py @@ -0,0 +1,743 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, exp2 +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32]) +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': 128, 'BV': 128}, num_warps=8, num_stages=3), + triton.Config({'BK': 64, 'BV': 64}, num_warps=4, num_stages=3), + triton.Config({'BK': 32, 'BV': 32}, num_warps=2, num_stages=3), + ], + key=['H', 'K', 'V', 'BT', 'TRANSPOSE_STATE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + g_gamma, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K*V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + + # [BT, BK] @ [BK, BV] -> [BT, BV] + if TRANSPOSE_STATE: + b_o += tl.dot(b_q, tl.trans(b_h)) + else: + b_o += tl.dot(b_q, b_h) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + if USE_EXP2: + b_o = b_o * exp2(b_g)[:, None] + b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) + else: + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + if USE_EXP2: + b_o = b_o * exp2(b_g)[:, None] + b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) + else: + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_v = tl.load(p_v, boundary_check=(0, 1)) + # to fix mma -> mma layout conversion + # already solved by triton v3.2 or higher + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'USE_DW': lambda args: args['dw'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dqkwg( + q, + k, + v, + g, + g_gamma, + h, + do, + dh, + dq, + dk, + dw, + dv, + dg, + cu_seqlens, + chunk_indices, + scale, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, + USE_DW: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + all = B * T + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K*V + dh += (i_tg * H + i_h).to(tl.int64) * K*V + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + + # for delta rule only + if USE_DW: + dw += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + + if USE_G: + dg += i_k * all * H + b_dg_last = tl.zeros([1], dtype=tl.float32) if USE_G else None + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_ds = tl.zeros([BT, BT], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + if USE_G: + b_dg_last += (tl.sum(b_h * b_dh)) + # [BT, BV] @ [BV, BT] -> [BT, BT] + b_ds += tl.dot(b_do, tl.trans(b_v)) + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + if USE_DW: + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) + + if USE_DW: + p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + if USE_G: + b_dg = tl.zeros([BT], dtype=tl.float32) + g += bos * H + i_h + dg += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + if USE_EXP2: + b_dg_last *= exp2(b_g_last) + b_dq = b_dq * exp2(b_g)[:, None] * scale + else: + b_dg_last *= exp(b_g_last) + b_dq = b_dq * exp(b_g)[:, None] * scale + b_dg += tl.sum(b_dq * b_q, axis=1) + + if USE_EXP2: + b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + else: + b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + b_dg -= tl.sum(b_k * b_dk, axis=1) + b_dg_last += tl.sum(b_dk * b_k) + + if USE_EXP2: + b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale + else: + b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale + b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k)) + b_dg += tl.sum(b_ds2, axis=1) + b_dg -= tl.sum(b_ds2, axis=0) + + b_ds = b_ds.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue + # b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last) + b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + elif USE_G_GAMMA: + if USE_EXP2: + b_dq = b_dq * exp2(b_g)[:, None] * scale + b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale + else: + b_dq = b_dq * exp(b_g)[:, None] * scale + b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale + b_ds = b_ds.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + else: + b_ds = tl.where(m_A, b_ds, 0) + b_ds = b_ds.to(b_k.dtype) + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) * scale + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dv( + q, + k, + g, + g_gamma, + do, + dv, + dh, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + dh += (i_tg * H + i_h).to(tl.int64) * K*V + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, b_q) + p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + b_g_last = b_gamma * min(BT, T - i_t * BT) + + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + if USE_G or USE_G_GAMMA: + if USE_EXP2: + b_A = tl.where(m_A, b_A * exp2(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) + b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] + else: + b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) + b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] + else: + b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv += tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'USE_A': lambda args: args['A'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dv_local( + q, + k, + g, + g_gamma, + A, + do, + dv, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_EXP2: tl.constexpr, + USE_A: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + + if USE_A: + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + else: + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (tl.arange(0, BT) + 1) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_A += tl.dot(b_k, b_q) * scale + if USE_G or USE_G_GAMMA: + if USE_EXP2: + b_A *= exp2(b_g[None, :] - b_g[:, None]) + else: + b_A *= exp(b_g[None, :] - b_g[:, None]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) + + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> torch.Tensor: + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o = torch.empty_like(v) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_fwd_kernel_o[grid]( + q=q, + k=k, + v=v, + h=h, + g=g, + g_gamma=g_gamma, + o=o, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + ) + return o + + +def chunk_bwd_dv( + q: torch.Tensor, + k: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> torch.Tensor: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # H100 can have larger block size + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem('ada', k.device.index): + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NV = triton.cdiv(V, BV) + if scale is None: + scale = k.shape[-1] ** -0.5 + + dv = torch.empty_like(do) + grid = (NV, NT, B * H) + chunk_bwd_kernel_dv[grid]( + q=q, + k=k, + g=g, + g_gamma=g_gamma, + do=do, + dv=dv, + dh=dh, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + # H200 Triton autotune can wedge while timing this backward kernel. + # Use the same conservative fixed launch across ranks. + num_warps=4, + num_stages=3, + ) + return dv + + +def chunk_bwd_dv_local( + q: torch.Tensor, + k: torch.Tensor, + do: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + A: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> torch.Tensor: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # H100 can have larger block size + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem('ada', k.device.index): + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_bwd_kernel_dv_local[grid]( + q=q, + k=k, + g=g, + g_gamma=g_gamma, + A=A, + do=do, + dv=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + # Keep the local DV variant deterministic too; otherwise it can hit + # the same H200 autotune stall after the global DV kernel is fixed. + num_warps=4, + num_stages=3, + ) + return dv + + +def chunk_bwd_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + w: torch.Tensor | None = None, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + dv: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem('ada', k.device.index): + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NK = triton.cdiv(K, BK) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None + dw = torch.empty_like(w) if w is not None else None + + grid = (NK, NT, B * H) + # H200 Triton autotune can stall for this backward kernel while benchmarking + # configs. Use a fixed conservative launch; avoiding autotune matters more + # than chasing a small per-kernel speedup during long 20B training runs. + chunk_bwd_kernel_dqkwg[grid]( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + h=h, + do=do, + dh=dh, + dw=dw, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + num_warps=4, + num_stages=3, + ) + + if dg is not None: + dg = dg.sum(0) + return dq, dk, dw, dg diff --git a/fla/ops/common/chunk_scaled_dot_kkt.py b/fla/ops/common/chunk_scaled_dot_kkt.py new file mode 100644 index 0000000000000000000000000000000000000000..2c33a12db4e9c82cb5ac68b60fb3912bba15f3ed --- /dev/null +++ b/fla/ops/common/chunk_scaled_dot_kkt.py @@ -0,0 +1,126 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A *= exp(b_g_diff) + b_A *= b_b[:, None] + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_kkt_fwd( + k: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + ) + return A diff --git a/fla/ops/common/fused_chunk.py b/fla/ops/common/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..e55e6a72751ff8c451a3293769fd948f1f52ea8c --- /dev/null +++ b/fla/ops/common/fused_chunk.py @@ -0,0 +1,635 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.op import exp +from fla.utils import ( + IS_NVIDIA_HOPPER, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_shared_mem, + input_guard, +) + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_chunk_fwd_kernel( + q, + k, + v, + g, + g_gamma, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + + all = B * T + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + + o_i = tl.arange(0, BT) + + if USE_G_GAMMA: + # decay rate given the head index + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (o_i + 1) + b_g_last = b_gamma * BT + b_gq = exp(b_g) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + + # [BT, BT] + m_s = o_i[:, None] >= o_i[None, :] + + q = q + (bos*H + i_h) * K + k = k + (bos*H + i_h) * K + v = v + (bos*H + i_h) * V + o = o + (i_k * all + bos).to(tl.int64) * H*V + i_h * V + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(0, NT): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + + # [BT, BT] + b_s = tl.dot(b_q, b_k) + + # scalar decay + if USE_G: + p_g = g + (bos + o_t) * H + i_h + b_g = tl.load(p_g, mask=(o_t < T), other=0.) + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + + b_gq = exp(b_g) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + if USE_G_GAMMA: + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + if USE_G or USE_G_GAMMA: + b_gs = tl.where(m_s & m_t, exp(b_g[:, None] - b_g[None, :]), 0) + # [BT, BT] + b_s *= b_gs + # [BT, BV] + b_o = tl.dot(b_s.to(b_q.dtype), b_v) + tl.dot(b_q, b_h.to(b_q.dtype)) * b_gq[:, None] + b_v = (b_v * b_gk[:, None]).to(b_v.dtype) + b_h *= b_gn + else: + # [BT, BT] + b_s *= m_s & m_t + # [BT, BV] + b_o = tl.dot(b_s.to(b_q.dtype), b_v) + tl.dot(b_q, b_h.to(b_q.dtype)) + + b_h += tl.dot(b_k, b_v) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE': lambda args: args['dht'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_chunk_bwd_kernel( + q, + k, + v, + g, + g_gamma, + do, + dq, + dk, + dv, + dg, + h0, + dht, + dh0, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + + all = B * T + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + NV = tl.cdiv(V, BV) + + o_i = tl.arange(0, BT) + if USE_G_GAMMA: + b_gamma = tl.load(g_gamma + i_h) + b_g = b_gamma * (o_i + 1) + b_g_last = b_gamma * BT + b_gq = exp(b_g) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + + m_s = o_i[:, None] >= o_i[None, :] + + q = q + (bos*H + i_h) * K + k = k + (bos*H + i_h) * K + v = v + (bos*H + i_h) * V + do = do + (bos*H + i_h) * V + dq = dq + (i_v * all + bos).to(tl.int64) * H*K + i_h * K + dk = dk + (i_v * all + bos).to(tl.int64) * H*K + i_h * K + dv = dv + (i_k * all + bos).to(tl.int64) * H*V + i_h * V + + # [BV, BK] + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h = tl.make_block_ptr(h0 + i_nh * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(0, NT): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BT] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + + # [BT, BT] + b_ds = tl.dot(b_do, b_v) * scale + + # scalar decay + if USE_G: + p_g = g + (bos + o_t) * H + i_h + b_g = tl.load(p_g, mask=(o_t < T), other=0.) + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + + b_gq = exp(b_g) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + + p_dg = dg + ((i_k * NV + i_v) * all + (bos + o_t)).to(tl.int64) * H + i_h + # [BT, BT] + b_gs = tl.where(m_s & m_t, exp(b_g[:, None] - b_g[None, :]), 0) + b_ds = b_ds * b_gs + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) + # [BT] + b_dg_t = tl.sum(b_q * b_dq, 1) + tl.store(p_dg, b_dg_t.to(p_dg.dtype.element_ty), mask=m_t) + # [BV, BK] + b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) + + elif USE_G_GAMMA: + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + + # [BT, BT] + b_gs = tl.where(m_s & m_t, exp(b_g[:, None] - b_g[None, :]), 0) + b_ds = b_ds * b_gs + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) + # [BV, BK] + b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) + + else: + # [BT, BT] + b_ds *= m_s & m_t + # [BT, BK] + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * scale).to(b_k.dtype), b_h.to(b_k.dtype)) + # [BV, BK] + b_h += tl.dot(b_v, b_k) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE: + p_dh = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) + + if USE_G: + b_dg = tl.zeros([BT], dtype=tl.float32) + b_dg_last = tl.sum(tl.trans(b_h) * b_dh) + + # sync threads + b_h = None + tl.debug_barrier() + + for i_t in range(NT - 1, -1, -1): + p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BK, BT] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + # [BT, BT] + b_s = tl.dot(b_k, b_q) + b_ds = tl.dot(b_v, tl.trans(b_do)) + + if USE_G: + p_g = g + (bos + o_t) * H + i_h + p_dg = dg + ((i_k * NV + i_v) * all + (bos + o_t)).to(tl.int64) * H + i_h + b_g = tl.load(p_g, mask=m_t, other=0.) + b_g_last = tl.load(g + (bos + last_idx) * H + i_h) + + b_gq = exp(b_g) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + b_gs = tl.trans(tl.where(m_s & (m_t[:, None] & m_t), exp(b_g[:, None] - b_g[None, :]), 0)) * scale + + b_s = b_s * b_gs + b_ds = b_ds * b_gs + + # [BT, BK] + b_dk = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) * b_gk[:, None] + + # [BT] + b_dg_t = tl.where(m_t, tl.load(p_dg, mask=m_t, other=0.) - tl.sum(b_k * b_dk, 1), 0) + b_dg_last += tl.sum(b_dg_t, 0) + b_dg = b_dg_last + b_dg_t - tl.cumsum(b_dg_t, 0) + + # [BT, BV] + b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) * b_gk[:, None] + # [BK, BV] + b_dh = b_dh * b_gn + tl.dot(b_q, (b_do * b_gq[:, None] * scale).to(b_do.dtype)) + + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + + elif USE_G_GAMMA: + b_g_last = b_gamma * min(BT, T - i_t * BT) + b_gk = exp(b_g_last - b_g) + b_gn = exp(b_g_last) + b_gs = tl.trans(tl.where(m_s & (m_t[:, None] & m_t), exp(b_g[:, None] - b_g[None, :]), 0)) * scale + + b_s = b_s * b_gs + b_ds = b_ds * b_gs + + b_dk = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) * b_gk[:, None] + # [BT, BV] + b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) * b_gk[:, None] + # [BK, BV] + b_dh = b_dh * b_gn + tl.dot(b_q, (b_do * b_gq[:, None] * scale).to(b_do.dtype)) + + else: + mask = tl.trans(m_s & (m_t[:, None] & m_t)) + b_s = tl.where(mask, b_s * scale, 0).to(b_do.dtype) + b_ds = tl.where(mask, b_ds * scale, 0).to(b_q.dtype) + + b_dk = tl.dot(b_ds, tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) + # [BT, BV] + b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) + # [BK, BV] + b_dh += tl.dot(b_q, (b_do * scale).to(b_do.dtype)) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_chunk_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + BK = min(max(triton.next_power_of_2(K), 16), 64) + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + NK = triton.cdiv(K, BK) + + o = v.new_empty(NK, *v.shape, dtype=torch.float) if NK > 1 else torch.empty_like(v) + ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + def grid(meta): return (triton.cdiv(V, meta['BV']), NK, N * H) + fused_chunk_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + o=o, + h0=initial_state, + ht=ht, + cu_seqlens=cu_seqlens, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + ) + if NK > 1: + o = o.sum(0).to(v) + return o, ht + + +def fused_chunk_bwd( + q, + k, + v, + g, + g_gamma, + do, + scale, + initial_state: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +): + B, T, H, K, V = *q.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BT = chunk_size + BK = min(max(triton.next_power_of_2(K), 16), 64) + BV = min(max(triton.next_power_of_2(V), 16), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + dq = q.new_empty(NV, *q.shape, dtype=torch.float) if NV > 1 else torch.empty_like(q) + dk = k.new_empty(NV, *k.shape, dtype=torch.float) if NV > 1 else torch.empty_like(k) + dv = v.new_empty(NK, *v.shape, dtype=torch.float) if NK > 1 else torch.empty_like(v) + dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float) if g is not None else None + dh0 = torch.empty_like(initial_state) if initial_state is not None else None + + grid = (NV, NK, N * H) + fused_chunk_bwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + do=do, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + h0=initial_state, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dq = dq.sum(0) if NV > 1 else dq + dk = dk.sum(0) if NV > 1 else dk + dv = dv.sum(0) if NK > 1 else dv + if dg is not None: + dg = dg.sum(0).to(g) + + return dq, dk, dv, dg, dh0 + + +class FusedChunkFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + g, + g_gamma, + scale, + initial_state, + output_final_state, + cu_seqlens, + ): + chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) if g is not None else None + o, ht = fused_chunk_fwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + + ctx.save_for_backward(q, k, v, g, g_gamma, initial_state) + ctx.chunk_size = chunk_size + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht=None): + q, k, v, g, g_gamma, initial_state = ctx.saved_tensors + + dq, dk, dv, dg, dh0 = fused_chunk_bwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + do=do, + scale=ctx.scale, + initial_state=initial_state, + dht=dht, + cu_seqlens=ctx.cu_seqlens, + chunk_size=ctx.chunk_size, + ) + if g is not None: + dg = dg.to(g) + return dq.to(q), dk.to(k), dv.to(v), dg, None, None, dh0, None, None + + +def fused_chunk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + Compared to GLA, the gating is head-wise instead of elementwise. + g_gamma (torch.Tensor): + Log decay of shape `[H]`. + Head-wise data-independent decay is used if `g_gamma` is provided. + Only one of `g` or `g_gamma` should be provided. + scale (Optional[int]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + if g is not None and g_gamma is not None: + raise ValueError("Only one of `g` or `g_gamma` should be provided.") + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = FusedChunkFunction.apply( + q, + k, + v, + g, + g_gamma, + scale, + initial_state, + output_final_state, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/common/fused_recurrent.py b/fla/ops/common/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..1990f872f522f2a784207cf643a085450e53f50f --- /dev/null +++ b/fla/ops/common/fused_recurrent.py @@ -0,0 +1,567 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [4, 8] + ], + key=['BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['B', 'T']) +def fused_recurrent_fwd_kernel( + q, + k, + v, + g, + g_gamma, + gk, + gv, + o, + h0, + ht, + cu_seqlens, + scale, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + + all = B * T + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_o = o + ((i_k * all + bos) + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + if USE_G: + p_g = g + (bos + ((T-1) if REVERSE else 0)) * H + i_h + if USE_GK: + p_gk = gk + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + if USE_GV: + p_gv = gv + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + if USE_G_GAMMA: + b_g_gamma = tl.load(g_gamma + i_h) + + m_k = o_k < K + m_v = o_v < V + m_h = m_k[:, None] & m_v[None, :] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=m_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=m_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + b_h = b_h * exp(b_g) + if USE_G_GAMMA: + b_h = b_h * exp(b_g_gamma) + if USE_GK: + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_h = b_h * exp(b_gk[:, None]) + if USE_GV: + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + b_h = b_h * exp(b_gv[None, :]) + b_h += b_k[:, None] * b_v[None, :] + b_o = b_h * b_q[:, None] + b_o = tl.sum(b_o, axis=0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) + p_q += (-1 if REVERSE else 1) * H*K + p_k += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_o += (-1 if REVERSE else 1) * H*V + if USE_G: + p_g += (-1 if REVERSE else 1) * H + if USE_GK: + p_gk += (-1 if REVERSE else 1) * H*K + if USE_GV: + p_gv += (-1 if REVERSE else 1) * H*V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_h) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [4] + ], + key=['BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_GK', 'USE_GV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['B', 'T']) +def fused_recurrent_bwd_kernel( + q, + k, + v, + g, + g_gamma, + gk, + gv, + o, + h0, + do, + dq, + dk, + dv, + dg, + dgk, + dgv, + dht, + dh0, + cu_seqlens, + scale, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_G: tl.constexpr, + USE_G_GAMMA: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + + all = B * T + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + NV = tl.cdiv(V, BV) + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_h = m_k[:, None] & m_v[None, :] + + p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_do = do + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_dq = dq + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + if USE_G: + p_g = g + (bos + ((T-1) if REVERSE else 0)) * H + i_h + if USE_GK: + p_gk = gk + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + if USE_GV: + p_gv = gv + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + if USE_G_GAMMA: + b_g_gamma = tl.load(g_gamma + i_h) + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=m_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=m_v, other=0).to(tl.float32) + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + b_h = b_h * exp(b_g) + if USE_G_GAMMA: + b_h = b_h * exp(b_g_gamma) + if USE_GK: + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_h = b_h * exp(b_gk[:, None]) + if USE_GV: + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + b_h = b_h * exp(b_gv[None, :]) + b_h += b_k[:, None] * b_v[None, :] + b_dq = b_h * b_do[None, :] + b_dq = tl.sum(b_dq, axis=1) * scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_k) + + p_k += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_do += (-1 if REVERSE else 1) * H*V + p_dq += (-1 if REVERSE else 1) * H*K + if USE_G: + p_g += (-1 if REVERSE else 1) * H + if USE_GK: + p_gk += (-1 if REVERSE else 1) * H*K + if USE_GV: + p_gv += (-1 if REVERSE else 1) * H*V + + # sync threads + tl.debug_barrier() + + p_q = q + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + + p_do = do + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_dq = dq + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_dk = dk + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_dv = dv + ((i_k * all + bos) + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + if USE_G: + p_g = g + (bos + ((T - 1) if not REVERSE else 0)) * H + i_h + p_dg = dg + ((i_k * NV + i_v) * all + bos + ((T - 1) if not REVERSE else 0)) * H + i_h + if USE_GK: + p_gk = gk + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_dgk = dgk + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k + if USE_GV: + p_o = o + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_gv = gv + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_dgv = dgv + ((i_k * all + bos) + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_h, other=0).to(tl.float32) + + if USE_G: + b_dg = tl.sum(b_h * b_dh) + if USE_GK: + b_dgk = tl.sum(b_h * b_dh, 1) + if USE_GV: + b_dgv = tl.sum(b_h * b_dh, 0) + + for _ in range(T): + b_q = tl.load(p_q, mask=m_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=m_v, other=0).to(tl.float32) + b_dh += (b_q * scale)[:, None] * b_do[None, :] + b_dk = tl.sum(b_dh * b_v[None, :], axis=1) + b_dv = tl.sum(b_dh * b_k[:, None], axis=0) + + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + b_dq = tl.load(p_dq, mask=m_k, other=0).to(tl.float32) + b_dg += tl.sum(b_q * b_dq - b_k * b_dk) + b_dh *= exp(b_g) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty)) + if USE_G_GAMMA: + b_dh *= exp(b_g_gamma) + if USE_GK: + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_dq = tl.load(p_dq, mask=m_k, other=0).to(tl.float32) + b_dgk += b_q * b_dq - b_k * b_dk + b_dh *= exp(b_gk)[:, None] + tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), mask=m_k) + if USE_GV: + b_o = tl.load(p_o, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + if i_k == 0: + b_dgv += b_o * b_do + b_dgv -= b_v * b_dv + b_dh *= exp(b_gv)[None, :] + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), mask=m_v) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) + + p_q += (1 if REVERSE else -1) * H*K + p_k += (1 if REVERSE else -1) * H*K + p_v += (1 if REVERSE else -1) * H*V + + p_do += (1 if REVERSE else -1) * H*V + p_dq += (1 if REVERSE else -1) * H*K + p_dk += (1 if REVERSE else -1) * H*K + p_dv += (1 if REVERSE else -1) * H*V + if USE_G: + p_g += (1 if REVERSE else -1) * H + p_dg += (1 if REVERSE else -1) * H + if USE_GK: + p_gk += (1 if REVERSE else -1) * H*K + p_dgk += (1 if REVERSE else -1) * H*K + if USE_GV: + p_o += (1 if REVERSE else -1) * H*V + p_gv += (1 if REVERSE else -1) * H*V + p_dgv += (1 if REVERSE else -1) * H*V + + if STORE_INITIAL_STATE_GRADIENT: + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) + + +def fused_recurrent_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + h0 = initial_state + ht = q.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + o = q.new_empty(NK, *v.shape, dtype=torch.float32) + + grid = (NV, NK, N * H) + fused_recurrent_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + o=o, + h0=h0, + ht=ht, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + USE_G=g is not None, + USE_G_GAMMA=g_gamma is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + REVERSE=reverse, + ) + o = o.sum(0) + return o, ht + + +def fused_recurrent_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + o: torch.Tensor | None = None, + do: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + h0 = initial_state + dq = q.new_empty(NV, *q.shape, dtype=torch.float32) + dk = q.new_empty(NV, *k.shape, dtype=torch.float32) + dv = q.new_empty(NK, *v.shape, dtype=torch.float32) + dh0 = torch.empty_like(h0) if h0 is not None else None + + dg, dgk, dgv = None, None, None + if g is not None: + dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float32) + if gk is not None: + dgk = gk.new_empty(NV, *gk.shape, dtype=torch.float32) + if gv is not None: + dgv = gv.new_empty(NK, *gv.shape, dtype=torch.float32) + + grid = (NV, NK, N * H) + fused_recurrent_bwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + o=o, + h0=h0, + do=do, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + dgk=dgk, + dgv=dgv, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + USE_G=g is not None, + USE_G_GAMMA=g_gamma is not None, + USE_GK=gk is not None, + USE_GV=gv is not None, + REVERSE=reverse, + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + if g is not None: + dg = dg.sum(0).to(g) + if gk is not None: + dgk = dgk.sum(0).to(gk) + if gv is not None: + dgv = dgv.sum(0).to(gv) + + return dq, dk, dv, dg, dgk, dgv, dh0 + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, ht = fused_recurrent_fwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, k, v, g, g_gamma, gk, gv, initial_state, o) + ctx.scale = scale + ctx.reverse = reverse + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + q, k, v, g, g_gamma, gk, gv, initial_state, o = ctx.saved_tensors + dq, dk, dv, dg, dgk, dgv, dh0 = fused_recurrent_bwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=gk, + gv=gv, + o=o, + do=do, + dht=dht, + scale=ctx.scale, + initial_state=initial_state, + reverse=ctx.reverse, + cu_seqlens=ctx.cu_seqlens, + ) + return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dg, None, dgk, dgv, None, dh0, None, None, None + + +def fused_recurrent( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + return FusedRecurrentFunction.apply( + q, + k, + v, + g, + g_gamma, + gk, + gv, + scale, + initial_state, + output_final_state, + reverse, + cu_seqlens, + ) diff --git a/fla/ops/common/intracard_cp.py b/fla/ops/common/intracard_cp.py new file mode 100644 index 0000000000000000000000000000000000000000..e73a308f96bc32ac401d28b1233a3ad096e03e44 --- /dev/null +++ b/fla/ops/common/intracard_cp.py @@ -0,0 +1,597 @@ +"""Intra-Card Context Parallel for KDA inference (varlen mode only). + +Optimized: all CPU-side index computation uses pure Python loops instead of +torch tensor operations (repeat_interleave, arange, cumsum, etc.) to eliminate +per-op overhead on tiny arrays. GPU tensors are created directly from Python +lists to minimize cudaStreamSynchronize calls. +""" + +from __future__ import annotations + +import logging +import weakref +from collections import OrderedDict +from typing import NamedTuple + +import torch +import triton + +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_kernel_h_blockdim64 +from fla.ops.cp.chunk_delta_h import pre_process_fwd_kernel_merged +from fla.ops.utils.index import prepare_chunk_indices, prepare_chunk_offsets +from fla.utils import get_multiprocessor_count + +logger = logging.getLogger(__name__) + + +# Cache for intracard_fwd_h precomputation (Python results + GPU tensors) +# Key: object id of cu_seqlens (consistent with tensor_cache philosophy) +_intracard_cache: OrderedDict[tuple, _CacheEntry] = OrderedDict() +_INTRACARD_CACHE_MAXSIZE = 32 + + +class _CacheEntry(NamedTuple): + """Cache entry for intracard_fwd_h precomputation. + + Caches both Python computation results and GPU tensors to eliminate + redundant CPU→GPU transfers and Python loop computation. + """ + # Keep a weak reference to validate id-based key safety. + # If Python reuses an object id after GC, this guard prevents stale hits. + cu_seqlens_ref: weakref.ReferenceType[torch.Tensor] + # From prepare_subseq_cu_seqlens + cu_seqlens_subseq_values: list[int] + split_info: SplitSeqInfo + total_subseqs: int + # From _precompute_intracard_indices + cu_seqlens_split_values: list[int] + S_split_total: int + non_first_indices: list[int] + first_subseq_indices: list[int] + last_subseq_indices: list[int] + num_non_first: int + merge_seq_offsets: list[int] + merge_init_offsets: list[int] + # GPU tensors (cached to avoid H2D transfer) + cu_seqlens_subseq_gpu: torch.Tensor + cu_seqlens_split_flat: torch.Tensor + + +class SplitSeqInfo(NamedTuple): + """Information about split sequences (Python lists for zero-overhead access).""" + split_seq_ids: list[int] # [num_split_seqs] original sequence indices + start_subseq_idx: list[int] # [num_split_seqs] start index in subseq array + num_subseqs: list[int] # [num_split_seqs] number of sub-sequences per split + + @property + def num_split_seqs(self) -> int: + return len(self.split_seq_ids) + + def __bool__(self) -> bool: + return self.num_split_seqs > 0 + + +def _raw_chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V = *k.shape, u.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + if transpose_state_layout: + h = k.new_empty(B, NT, H, V, K) + final_state = k.new_zeros(N, H, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_zeros(N, H, K, V, dtype=torch.float32) if output_final_state else None + v_new = torch.empty_like(u) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta['BV']), N * H) + + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, v=u, w=w, v_new=v_new, + g=g, gk=gk, h=h, h0=initial_state, ht=final_state, + cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + T=T, H=H, K=K, V=V, BT=BT, USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + ) + return h, v_new, final_state + + +def compute_subseq_len( + seq_len: int, + num_sms: int, + num_heads: int, + chunk_size: int = 64, +) -> int: + """Compute sub-sequence length for intracard splitting. + + For linear recurrence (fwd_h), the sequential scan is the bottleneck. + Splitting always reduces the critical path and helps, as long as the + sequence is long enough to amortize the pre_scan + merge overhead. + + The fwd_h kernel grid is (num_v_blocks, N*H) where num_v_blocks ≈ 2. + Each sub-sequence contributes 2*H blocks. We target enough splits so + that even a single long sequence can saturate all SMs. + + A floor on subseq_chunks (MIN_SUBSEQ_CHUNKS) prevents subseq_len from + being too small, which would cause prepare_subseq_cu_seqlens to + unnecessarily split shorter sequences in mixed-length batches + (split threshold = 2 * subseq_len). + """ + seq_chunks = (seq_len + chunk_size - 1) // chunk_size + + if seq_chunks < 8: + return seq_len + + # Target splits: saturate SMs with the longest sequence alone. + # Each sub-seq contributes NUM_V_BLOCKS * num_heads blocks. + # Always at least 4 — for linear recurrence, CP4 always helps. + NUM_V_BLOCKS = 2 + target_splits = max(4, num_sms // (NUM_V_BLOCKS * num_heads)) + + subseq_chunks = (seq_chunks + target_splits - 1) // target_splits + + # Floor: prevent subseq_len from being too small. + # With chunk_size=64, MIN_SUBSEQ_CHUNKS=128 → subseq_len >= 8192 tokens, + # split threshold (3 * subseq_len) = 24576 tokens. + # Sequences shorter than it won't be split. + MIN_SUBSEQ_CHUNKS = 128 + subseq_chunks = max(subseq_chunks, MIN_SUBSEQ_CHUNKS) + + return subseq_chunks * chunk_size + + +def prepare_subseq_cu_seqlens( + cu_seqlens_cpu: torch.Tensor, + subseq_len: int, + chunk_size: int = 64, + max_splits: int = 32, +) -> tuple[list[int], SplitSeqInfo | bool, int]: + """Insert subseq split points into original cu_seqlens. + + Optimized: uses pure Python loops instead of torch tensor operations + for the small index arrays (typically 1-32 elements). + + Returns: + boundaries: List of cu_seqlens boundaries (can be used directly by _precompute_intracard_indices) + split_info: SplitSeqInfo for sequences that need splitting, or False if no splitting needed + total_subseqs: Total number of sub-sequences after splitting + """ + N = len(cu_seqlens_cpu) - 1 + if N == 0: + return cu_seqlens_cpu.tolist(), False, 0 + + subseq_chunks = (subseq_len + chunk_size - 1) // chunk_size + threshold_subseq_len = 3 * subseq_len + + split_seq_ids: list[int] = [] + start_subseq_idxs: list[int] = [] + num_subseqs_list: list[int] = [] + + # Build boundaries using pure Python loop + boundaries: list[int] = [0] + cumsum_offset = 0 + + for i in range(N): + seq_start = int(cu_seqlens_cpu[i].item()) + seq_end = int(cu_seqlens_cpu[i + 1].item()) + seq_len_i = seq_end - seq_start + seq_chunks_i = (seq_len_i + chunk_size - 1) // chunk_size + + if seq_len_i >= threshold_subseq_len: + # This sequence needs splitting + num_ss = min(max_splits, (seq_chunks_i + subseq_chunks - 1) // subseq_chunks) + chunks_per = (seq_chunks_i + num_ss - 1) // num_ss + actual_ssl = chunks_per * chunk_size + + split_seq_ids.append(i) + start_subseq_idxs.append(cumsum_offset) + num_subseqs_list.append(num_ss) + + for j in range(num_ss): + boundary = min(seq_start + (j + 1) * actual_ssl, seq_end) + boundaries.append(boundary) + cumsum_offset += num_ss + else: + # No split needed, single sub-sequence + boundaries.append(seq_end) + cumsum_offset += 1 + + if not split_seq_ids: + return cu_seqlens_cpu.tolist(), False, 0 + + total_subseqs = cumsum_offset + + split_info = SplitSeqInfo( + split_seq_ids=split_seq_ids, + start_subseq_idx=start_subseq_idxs, + num_subseqs=num_subseqs_list, + ) + + return boundaries, split_info, total_subseqs + + +def intracard_pre_scan( + kg: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + gk: torch.Tensor, + cu_seqlens_subseq_split: torch.Tensor, + S_split: int, + chunk_size: int = 64, + use_exp2: bool = True, +): + H, K, V = kg.shape[2], kg.shape[3], u.shape[3] + BK = triton.next_power_of_2(K) + BLOCK_SIZE = 32 if K <= 64 else 64 + + hm = kg.new_empty(S_split, H, K, V + K, dtype=torch.float32) + + grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H, S_split) + pre_process_fwd_kernel_merged[grid]( + k=kg, + v=u, + w=w, + g=None, + gk=gk, + hm=hm, + cu_seqlens=cu_seqlens_subseq_split, + T=0, + H=H, + K=K, + V=V, + BT=chunk_size, + BLOCK_SIZE=BLOCK_SIZE, + BK1=BK, + USE_EXP2=use_exp2, + MULTI_SEQS=True, + ) + + return hm + + +def intracard_merge( + hm: torch.Tensor, + split_info: SplitSeqInfo, + num_non_first: int, + merge_seq_offsets: list[int], + merge_init_offsets: list[int], + device: torch.device, + initial_state: torch.Tensor | None = None, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor | None, int]: + """Merge sub-sequence states using pre-computed parameters. + + All CPU-side preparation (cumsum, offset lists) is done in the caller + using pure Python loops. This function only creates GPU tensors and + launches the merge kernel. + """ + from fla.ops.cp.chunk_delta_h import merge_fwd_bwd_kernel + + if num_non_first == 0: + return None, 0 + + H = hm.shape[1] + K = hm.shape[2] + V = hm.shape[3] - K + BK = triton.next_power_of_2(K) + + num_split_seqs = split_info.num_split_seqs + + # Create all small GPU tensors from Python lists in one batch + # Merge into a single CPU→GPU transfer to minimize cudaStreamSynchronize + all_int_data = merge_seq_offsets + merge_init_offsets + split_info.split_seq_ids + all_tensor = torch.tensor(all_int_data, dtype=torch.int32, device=device) + n_so = len(merge_seq_offsets) + n_io = len(merge_init_offsets) + seq_offsets = all_tensor[:n_so] + init_offsets = all_tensor[n_so:n_so + n_io] + h0_seq_ids = all_tensor[n_so + n_io:] + + if transpose_state_layout: + initial_states_merge = hm.new_empty(num_non_first, H, V, K, dtype=torch.float32) + else: + initial_states_merge = hm.new_empty(num_non_first, H, K, V, dtype=torch.float32) + + def grid(meta): + return (triton.cdiv(V, meta['BV']), num_split_seqs, H) + + merge_fwd_bwd_kernel[grid]( + h=initial_states_merge, + ag_hm=hm, + pre_or_post_num_ranks=num_split_seqs, + rank=0, + seq_offsets=seq_offsets, + init_offsets=init_offsets, + h0_seq_ids=h0_seq_ids, + h0=initial_state, + H=H, + K=K, + V=V, + BK=BK, + FORWARD=True, + INTRACARD_MODE=True, + NUM_SEQ_ENTRIES=num_split_seqs, + TRANSPOSE_STATE=transpose_state_layout, + ) + + return initial_states_merge, num_non_first + + +def _precompute_intracard_indices( + split_info: SplitSeqInfo, + cu_seqlens_subseq_values: list[int], + N_orig: int, +) -> tuple[list[int], int, list[int], list[int], list[int], int, list[int], list[int]]: + """Pre-compute all derived indices using pure Python loops. + + Returns: + cu_seqlens_split_values: flattened cu_seqlens boundaries for split seqs (for pre_scan) + S_split_total: total number of sub-sequences from splits + non_first_indices: indices for scattering merge results into initial_state_expanded + first_subseq_indices: indices of first sub-sequence for each original sequence + last_subseq_indices: indices of last sub-sequence for each original sequence + num_non_first: total non-first sub-sequences (merge work) + merge_seq_offsets: cumulative sub-sequence counts for merge kernel + merge_init_offsets: cumulative non-first counts for merge kernel + """ + starts = split_info.start_subseq_idx + num_ss = split_info.num_subseqs + split_ids = split_info.split_seq_ids + + # cu_seqlens_split_values: for each split seq, extract [start:start+n+1] boundaries + cu_seqlens_split_values: list[int] = [] + S_split_total = 0 + for s, n in zip(starts, num_ss): + cu_seqlens_split_values.extend(cu_seqlens_subseq_values[s:s + n + 1]) + S_split_total += n + + # num_subseqs_per_seq: [N_orig], default 1 for unsplit sequences + num_subseqs_per_seq = [1] * N_orig + for sid, nss in zip(split_ids, num_ss): + num_subseqs_per_seq[sid] = nss + + # non_first_indices: for scattering merged initial states + non_first_indices: list[int] = [] + for s, n in zip(starts, num_ss): + for j in range(1, n): + non_first_indices.append(s + j) + + # first_subseq_indices: for scattering original initial states + first_subseq_indices: list[int] = [0] + running = 0 + for i in range(N_orig - 1): + running += num_subseqs_per_seq[i] + first_subseq_indices.append(running) + + # last_subseq_indices: for gathering final states + last_subseq_indices: list[int] = [] + running = 0 + for n in num_subseqs_per_seq: + running += n + last_subseq_indices.append(running - 1) + + # merge parameters + merge_seq_offsets: list[int] = [0] + merge_init_offsets: list[int] = [0] + for n in num_ss: + merge_seq_offsets.append(merge_seq_offsets[-1] + n) + merge_init_offsets.append(merge_init_offsets[-1] + n - 1) + num_non_first = merge_init_offsets[-1] + + return ( + cu_seqlens_split_values, + S_split_total, + non_first_indices, + first_subseq_indices, + last_subseq_indices, + num_non_first, + merge_seq_offsets, + merge_init_offsets, + ) + + +def intracard_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + max_splits: int = 32, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + assert cu_seqlens is not None, "intracard_fwd_h requires cu_seqlens" + + _, _, H, K, V = *k.shape, u.shape[-1] + device = k.device + + if cu_seqlens_cpu is None: + cu_seqlens_cpu = cu_seqlens.cpu() + + seq_lens = torch.diff(cu_seqlens_cpu) + max_seq_len = int(seq_lens.max().item()) + num_sms = get_multiprocessor_count() + subseq_len = compute_subseq_len(max_seq_len, num_sms, H, chunk_size) + + early_return = (seq_lens < 2 * subseq_len).all() + + cached = None + cache_key = None + + if not early_return: + # Use object identity (id) for cache key, consistent with tensor_cache philosophy + # vLLM slice creates new Python objects per batch, so id(cu_seqlens) is safe + cache_key = ( + id(cu_seqlens), # Object identity, not content hash + subseq_len, + chunk_size, + max_splits, + str(device), + ) + cached = _intracard_cache.get(cache_key) + if cached is not None: + # Guard against rare Python id reuse after original tensor is GC-ed. + # We only consider it a hit when the weakref points to the current object. + if cached.cu_seqlens_ref() is cu_seqlens: + _intracard_cache.move_to_end(cache_key) + else: + _intracard_cache.pop(cache_key, None) + cached = None + + if cached is not None: + # Cache hit: reuse all precomputed results including GPU tensors + cu_seqlens_subseq_values = cached.cu_seqlens_subseq_values + split_info = cached.split_info + total_subseqs = cached.total_subseqs + cu_seqlens_split_values = cached.cu_seqlens_split_values + S_split_total = cached.S_split_total + non_first_indices = cached.non_first_indices + first_subseq_indices = cached.first_subseq_indices + last_subseq_indices = cached.last_subseq_indices + num_non_first = cached.num_non_first + merge_seq_offsets = cached.merge_seq_offsets + merge_init_offsets = cached.merge_init_offsets + cu_seqlens_subseq_gpu = cached.cu_seqlens_subseq_gpu + cu_seqlens_split_flat = cached.cu_seqlens_split_flat + else: + # Cache miss: compute Python lists + cu_seqlens_subseq_values, split_info, total_subseqs = prepare_subseq_cu_seqlens( + cu_seqlens_cpu, subseq_len, chunk_size, max_splits=max_splits + ) + + if early_return or not split_info: + return _raw_chunk_gated_delta_rule_fwd_h( + k=k, w=w, u=u, g=g, gk=gk, + initial_state=initial_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + N_orig = len(cu_seqlens_cpu) - 1 + + if cached is None: + # Cache miss: continue Python computation and create GPU tensors + ( + cu_seqlens_split_values, + S_split_total, + non_first_indices, + first_subseq_indices, + last_subseq_indices, + num_non_first, + merge_seq_offsets, + merge_init_offsets, + ) = _precompute_intracard_indices(split_info, cu_seqlens_subseq_values, N_orig) + + # Create GPU tensors (will be cached for reuse) + dtype = cu_seqlens_cpu.dtype + cu_seqlens_subseq_gpu = torch.tensor(cu_seqlens_subseq_values, dtype=dtype, device=device) + cu_seqlens_split_flat = torch.tensor(cu_seqlens_split_values, dtype=dtype, device=device) + + # Store all results in cache (including GPU tensors to avoid H2D) + _intracard_cache[cache_key] = _CacheEntry( + cu_seqlens_ref=weakref.ref(cu_seqlens), + cu_seqlens_subseq_values=cu_seqlens_subseq_values, + split_info=split_info, + total_subseqs=total_subseqs, + cu_seqlens_split_values=cu_seqlens_split_values, + S_split_total=S_split_total, + non_first_indices=non_first_indices, + first_subseq_indices=first_subseq_indices, + last_subseq_indices=last_subseq_indices, + num_non_first=num_non_first, + merge_seq_offsets=merge_seq_offsets, + merge_init_offsets=merge_init_offsets, + cu_seqlens_subseq_gpu=cu_seqlens_subseq_gpu, + cu_seqlens_split_flat=cu_seqlens_split_flat, + ) + # Evict oldest entries if over capacity + while len(_intracard_cache) > _INTRACARD_CACHE_MAXSIZE: + _intracard_cache.popitem(last=False) + + hm = intracard_pre_scan( + kg=k, w=w, u=u, gk=gk, + cu_seqlens_subseq_split=cu_seqlens_split_flat, + S_split=S_split_total, + chunk_size=chunk_size, + use_exp2=use_exp2, + ) + + initial_states_merge, num_non_first = intracard_merge( + hm=hm, + split_info=split_info, + num_non_first=num_non_first, + merge_seq_offsets=merge_seq_offsets, + merge_init_offsets=merge_init_offsets, + device=device, + initial_state=initial_state, + transpose_state_layout=transpose_state_layout, + ) + + if transpose_state_layout: + initial_state_expanded = k.new_zeros(total_subseqs, H, V, K, dtype=torch.float32) + else: + initial_state_expanded = k.new_zeros(total_subseqs, H, K, V, dtype=torch.float32) + + if initial_state is not None: + initial_state_expanded[first_subseq_indices] = initial_state + + if initial_states_merge is not None and num_non_first > 0: + initial_state_expanded[non_first_indices] = initial_states_merge + + chunk_indices_subseq = prepare_chunk_indices(cu_seqlens_subseq_gpu, chunk_size) + + h, v_new, final_state_subseq = _raw_chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + gk=gk, + initial_state=initial_state_expanded, + output_final_state=output_final_state, + chunk_size=chunk_size, + save_new_value=save_new_value, + cu_seqlens=cu_seqlens_subseq_gpu, + chunk_indices=chunk_indices_subseq, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + if output_final_state and final_state_subseq is not None: + final_state = final_state_subseq[last_subseq_indices] + else: + final_state = final_state_subseq + + return h, v_new, final_state diff --git a/fla/ops/cp/KCP.md b/fla/ops/cp/KCP.md new file mode 100644 index 0000000000000000000000000000000000000000..7d35a807b076548360283703f66faed5944bf7ab --- /dev/null +++ b/fla/ops/cp/KCP.md @@ -0,0 +1,272 @@ +# KCP: Kimi Context Parallel + +Context Parallel for GDN (Gated Delta Rule) and KDA (Kimi Delta Attention). + +> CP was first introduced in [PR #691](https://github.com/fla-org/flash-linear-attention/pull/691). Special thanks to [mdy666](https://github.com/mdy666). + +## Notation + +Following the [Kimi Linear technical report](https://yzhang.site/assets/pubs/techreport/2025/kda.pdf) (Section 2.1): + +- **Vectors**: $\square_t \in \mathbb{R}^{d_k}$ or $\mathbb{R}^{d_v}$ for $\square \in \{\mathbf{q}, \mathbf{k}, \mathbf{v}, \mathbf{o}, \mathbf{u}, \mathbf{w}\}$ denotes the $t$-th column vector. +- **State**: $\mathbf{S}_t \in \mathbb{R}^{d_k \times d_v}$ is the matrix-form memory state. FLA kernels store it as $[d_k, d_v]$; some other backends transpose to $[d_v, d_k]$. +- **Chunk indexing**: The sequence of length $L$ is split into $L/C$ chunks of size $C$. $\square_{[t]} \in \mathbb{R}^{C \times d}$ stacks vectors within chunk $t$; $\square^r_{[t]} = \square_{tC+r}$ is the $r$-th element ($t \in [0, L/C)$, $r \in [1, C]$). State: $\mathbf{S}_{[t]} := \mathbf{S}^0_{[t]} = \mathbf{S}^C_{[t-1]}$. +- **Decay**: $\alpha_t \in [0,1]$ (GDN, scalar) or $\alpha_t \in [0,1]^{d_k}$ (KDA, per-dim). Cumulative decay $\gamma^{i \to j}_{[t]} := \prod_{k=i}^j \alpha^k_{[t]}$, abbreviated $\gamma^r_{[t]} := \gamma^{1 \to r}_{[t]}$. $\text{Diag}(\gamma^{i \to j}_{[t]}) := \prod_{k=i}^j \text{Diag}(\alpha^k_{[t]})$. $\Gamma^{i \to j}_{[t]} \in \mathbb{R}^{C \times d_k}$ stacks rows from $\gamma^i_{[t]}$ to $\gamma^j_{[t]}$. +- **Code mapping**: In code, `g` stores $\log(\alpha)$ (or $\log_2(\alpha)$ for KDA). After `chunk_local_cumsum`, `g` at position $r$ equals $\log \gamma^r_{[t]}$. Then $\exp(\texttt{g}) = \gamma$ and $\exp(\texttt{g\_last} - \texttt{g}_r) = \gamma^{r \to C}_{[t]}$. + +--- + +## Recurrence + +**GDN** — scalar per-head gate (Eq. from [Yang et al., 2025]): + +$$\mathbf{S}_t = \alpha_t (\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) \mathbf{S}_{t-1} + \beta_t \mathbf{k}_t \mathbf{v}_t^\top, \quad \mathbf{o}_t = \mathbf{S}_t^\top \mathbf{q}_t$$ + +**KDA** — per-dim gate (Eq. 1 in the report): + +$$\mathbf{S}_t = (\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) \, \text{Diag}(\alpha_t) \, \mathbf{S}_{t-1} + \beta_t \mathbf{k}_t \mathbf{v}_t^\top, \quad \mathbf{o}_t = \mathbf{S}_t^\top \mathbf{q}_t$$ + +In the chunkwise formulation, the WY representation (Eq. 7) computes auxiliary matrices $\mathbf{W}_{[t]}$ and $\mathbf{U}_{[t]}$. The inter-chunk state recurrence (Eq. 8) is: + +$$\mathbf{S}_{[t+1]} = \text{Diag}(\gamma^C_{[t]}) \, \mathbf{S}_{[t]} + \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right)^\top \left(\mathbf{U}_{[t]} - \mathbf{W}_{[t]} \, \mathbf{S}_{[t]}\right)$$ + +--- + +## GDN vs KDA: Gate Handling + +### GDN: scalar per-head gate + +- $\alpha_t \in [0,1]$, one scalar per head per token +- In code: `g` shape `[B, T, H]` where $\alpha = \exp(g)$; processed by `chunk_local_cumsum` +- All kernels receive **original** $\mathbf{k}$, $\mathbf{q}$, and **scalar** `g` +- Kernels internally apply gating via `USE_G=True`: + - Inter-chunk decay: $\mathbf{S} \leftarrow \gamma^C_{[t]} \cdot \mathbf{S}$ (scalar broadcast) + - Gated key: $\tilde{\mathbf{k}}^r_{[t]} = \mathbf{k}^r_{[t]} \cdot \gamma^{r \to C}_{[t]}$ (done inside kernel) + - Gated query: $\tilde{\mathbf{q}}^r_{[t]} = \mathbf{q}^r_{[t]} \cdot \gamma^r_{[t]}$ (done inside kernel, backward only) + +### KDA: per-dim gate + +- $\alpha_t \in [0,1]^{d_k}$, one value per dimension per token +- In code: `g` shape `[B, T, H, K]` where $\alpha = \exp_2(g)$; processed by `kda_gate_chunk_cumsum` (includes gate activation) or `chunk_local_cumsum` (if pre-computed) +- The WY representation step (`chunk_kda_fwd_intra` / `recompute_w_u_fwd`) pre-computes gated tensors: + - `kg`: row $r$ is $\mathbf{k}^r_{[t]} \odot \gamma^{r \to C}_{[t]}$, i.e., `k * exp2(gk_last - gk)`. In matrix form: $\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}$. + - `qg`: row $r$ is $\mathbf{q}^r_{[t]} \odot \gamma^{r}_{[t]}$, i.e., `q * exp2(gk)`. In matrix form: $\Gamma^{1 \to C}_{[t]} \odot \mathbf{Q}_{[t]}$ (saved for backward). +- All kernels receive **pre-gated** `kg` (and `qg` in backward), plus `gk=g` for inter-chunk decay +- Kernels apply only the **chunk-level** decay via `USE_GK=True`: + - Inter-chunk decay: $\mathbf{S} \leftarrow \text{Diag}(\gamma^C_{[t]}) \, \mathbf{S}$ (per-dim diagonal) + - No further gating on $\mathbf{k}$/$\mathbf{q}$ — already done externally + +**Why the difference**: GDN's scalar gate is cheap to apply inside kernels. KDA's per-dim gate $\text{Diag}(\alpha_t) \in \mathbb{R}^{d_k \times d_k}$ is more efficiently pre-applied during the WY representation step. + +--- + +## CP Architecture + +### Data Flow + +Each rank holds a local chunk of the sequence. CP computes cross-rank initial states via an all-gather + merge pattern: + +1. Each rank computes local $(\mathbf{S}_\text{ext}, \mathbf{M})$ from its chunk + - $\mathbf{S}_\text{ext} \in \mathbb{R}^{d_k \times d_v}$: accumulated state assuming $\mathbf{S}_0 = \mathbf{0}$ + - $\mathbf{M} \in \mathbb{R}^{d_k \times d_k}$: transition matrix (product of per-chunk transitions) + +2. All-gather $[\mathbf{S}_\text{ext}, \mathbf{M}]$ across all ranks + +3. Rank $r$ merges from ranks $< r$: + +$$\mathbf{S} = \mathbf{0}; \quad \text{for } j \text{ from } (r - n_\text{pre}) \text{ to } (r-1): \quad \mathbf{S} \leftarrow \mathbf{M}_j \, \mathbf{S} + \mathbf{S}_{\text{ext},j}$$ + +### Pre-Process Forward + +Computes $(\mathbf{S}_\text{ext}, \mathbf{M})$ for the local chunk. + +**Stage 1 — $\mathbf{S}_\text{ext} \in \mathbb{R}^{d_k \times d_v}$** (accumulated state): + +$$\mathbf{S} = \mathbf{0}$$ + +For each sub-chunk $[t]$: + +$$\mathbf{S} \leftarrow \text{Diag}(\gamma^C_{[t]}) \, \mathbf{S} + \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right)^\top \left(\mathbf{U}_{[t]} - \mathbf{W}_{[t]} \, \mathbf{S}\right)$$ + +**Stage 2 — $\mathbf{M} \in \mathbb{R}^{d_k \times d_k}$** (transition matrix): + +$$\mathbf{M} = \mathbf{I}$$ + +For each sub-chunk $[t]$: + +$$\mathbf{M}_{[t]} = \text{Diag}(\gamma^C_{[t]}) - \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right)^\top \mathbf{W}_{[t]}, \quad \mathbf{M} \leftarrow \mathbf{M}_{[t]} \, \mathbf{M}$$ + +**Merge (forward direction):** + +For rank $r$ with `pre_num_ranks` previous ranks: + +$$\mathbf{S} = \mathbf{0}; \quad \text{for } j \text{ from } (r - n_\text{pre}) \text{ to } (r-1): \quad \mathbf{S} \leftarrow \mathbf{M}_j \, \mathbf{S} + \mathbf{S}_{\text{ext},j}$$ + +### Pre-Process Backward + +Same structure but **reversed** direction — merges from ranks **after** current rank. + +**Stage 1 — $\mathrm{d}\mathbf{S}_\text{ext} \in \mathbb{R}^{d_k \times d_v}$:** + +$$\mathrm{d}\mathbf{S} = \mathbf{0}$$ + +For each sub-chunk $[t]$ (reverse order): + +$$\mathrm{d}\mathbf{S} \leftarrow \text{Diag}(\gamma^C_{[t]}) \, \mathrm{d}\mathbf{S}$$ +$$\mathrm{d}\mathbf{V} = \mathbf{K}_{[t]} \, \mathrm{d}\mathbf{S} + \mathrm{d}\mathbf{V}_\text{local}$$ +$$\mathrm{d}\mathbf{S} \leftarrow \mathrm{d}\mathbf{S} + \left(\Gamma^{1 \to C}_{[t]} \odot \mathbf{Q}_{[t]}\right)^\top \mathrm{d}\mathbf{O}_{[t]} \cdot s - \mathbf{W}_{[t]}^\top \, \mathrm{d}\mathbf{V}$$ + +where $s = d_k^{-1/2}$ is the scaling factor. + +**Stage 2 — $\mathrm{d}\mathbf{M} \in \mathbb{R}^{d_k \times d_k}$:** + +$$\mathrm{d}\mathbf{M} = \mathbf{I}$$ + +For each sub-chunk $[t]$ (reverse order): + +$$\mathrm{d}\mathbf{M}_{[t]} = \text{Diag}(\gamma^C_{[t]}) - \mathbf{W}_{[t]}^\top \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right)$$ +$$\mathrm{d}\mathbf{M} \leftarrow \mathrm{d}\mathbf{M}_{[t]} \, \mathrm{d}\mathbf{M}$$ + +Note: $\mathrm{d}\mathbf{M}_{[t]}$ is the transpose of the forward $\mathbf{M}_{[t]}$ ($\mathbf{W}^\top \mathbf{K}$ vs. $\mathbf{K}^\top \mathbf{W}$). + +**Merge (backward direction):** + +For rank $r$ with `post_num_ranks` following ranks: + +$$\mathrm{d}\mathbf{S} = \mathbf{0}; \quad \text{for } j \text{ from } (r + n_\text{post}) \text{ down to } (r+1): \quad \mathrm{d}\mathbf{S} \leftarrow \mathrm{d}\mathbf{M}_j \, \mathrm{d}\mathbf{S} + \mathrm{d}\mathbf{S}_{\text{ext},j}$$ + +--- + +## Actual Code Flow + +### GDN Forward + +```python +g = chunk_local_cumsum(g, chunk_size=64) +w, u = recompute_w_u_fwd(k, v, beta, A, g=g) + +# CP pre-process: original k, scalar g +initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, w=w, u=u, g=g, # USE_G=True, USE_GK=False + context=cp_context, +) + +# Main kernel: original k, scalar g +h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, w=w, u=u, g=g, + initial_state=initial_state, +) +``` + +### GDN Backward + +```python +w, u = recompute_w_u_fwd(k, v, beta, A, g=g) +h, v_new, _ = chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, ...) +dv = chunk_bwd_dv_local(q=q, k=k, g=g, do=do, ...) + +# CP pre-process: original q, k, scalar g +dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, k=k, w=w, do=do, dv=dv, g=g, # USE_G=True, USE_GK=False + context=cp_context, +) + +# Main kernel: original q, k, scalar g +dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, k=k, w=w, g=g, + dht=dht, ... +) +``` + +### KDA Forward + +```python +# 1. Intra-chunk: compute WY repr + pre-gated tensors +w, u, qg, kg, Aqk, Akk = chunk_kda_fwd_intra(q, k, v, gk=g, beta, ...) +# kg = K ⊙ exp2(γ^{r→C}_{[t]}) rows of Γ^{i→C} ⊙ K +# qg = Q ⊙ exp2(γ^r_{[t]}) rows of Γ^{1→C} ⊙ Q (saved for backward) + +# 2. CP pre-process: pre-gated kg, per-dim gk=g +initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=kg, w=w, u=u, gk=g, # USE_G=False, USE_GK=True, use_exp2=True + context=cp_context, +) + +# 3. Main kernel: pre-gated kg, per-dim gk=g +h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, w=w, u=u, gk=g, + initial_state=initial_state, + use_exp2=True, +) +``` + +### KDA Backward + +```python +# 1. Recompute WY repr +w, u, qg, kg = recompute_w_u_fwd(q, k, v, beta, A=Akk, gk=g, ...) +# qg = Q ⊙ exp2(γ^r_{[t]}) +# kg = K ⊙ exp2(γ^{r→C}_{[t]}) + +# 2. Recompute state +h, v_new, _ = chunk_gated_delta_rule_fwd_h(k=kg, w=w, u=u, gk=g, ...) + +# 3. Compute local dv +dAqk, dv = chunk_kda_bwd_dAv(q, k, v=v_new, do, A=Aqk, ...) + +# 4. CP pre-process: pre-gated qg, kg, per-dim gk=g +dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=qg, k=kg, w=w, do=do, dv=dv, gk=g, # USE_G=False, USE_GK=True, use_exp2=True + context=cp_context, +) + +# 5. Main kernel: pre-gated qg, kg +dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, k=kg, w=w, gk=g, + dht=dht, ... + use_exp2=True, +) +``` + +--- + +## Input Tensor Summary + +| Function | GDN | KDA | Gate Path | +|----------|-----|-----|-----------| +| pre_process_fwd | k=$\mathbf{k}$, g=g | k=`kg`, gk=g | GDN: `USE_G`, KDA: `USE_GK` | +| fwd_h | k=$\mathbf{k}$, g=g | k=`kg`, gk=g | Same as pre_process | +| pre_process_bwd | q=$\mathbf{q}$, k=$\mathbf{k}$, g=g | q=`qg`, k=`kg`, gk=g | GDN: `USE_G`, KDA: `USE_GK` | +| bwd_dhu | q=$\mathbf{q}$, k=$\mathbf{k}$, g=g | q=`qg`, k=`kg`, gk=g | Same as pre_process | + +**Key consistency**: pre_process and main kernel always receive the **same** tensors. +For KDA, both receive pre-gated `kg` ($= \Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}$) and `qg` ($= \Gamma^{1 \to C}_{[t]} \odot \mathbf{Q}_{[t]}$). +For GDN, both receive original $\mathbf{k}$, $\mathbf{q}$ (gating applied inside the kernel). + +--- + +## Transition Matrix $\mathbf{M}$ + +The transition matrix captures how the state transforms across a chunk. Derived from Eq. 8: + +$$\mathbf{M}_{[t]} = \text{Diag}(\gamma^C_{[t]}) - \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right)^\top \mathbf{W}_{[t]} \quad \text{(forward)}$$ + +$$\mathrm{d}\mathbf{M}_{[t]} = \text{Diag}(\gamma^C_{[t]}) - \mathbf{W}_{[t]}^\top \left(\Gamma^{i \to C}_{[t]} \odot \mathbf{K}_{[t]}\right) \quad \text{(backward, transposed)}$$ + +Where $\text{Diag}(\gamma^C_{[t]})$ is: +- GDN: $\exp(g_\text{last}) \cdot \mathbf{I}$ — scalar times identity +- KDA: $\text{Diag}(\gamma^C_{[t]})$ — per-dim diagonal, where $\gamma^C_{[t]} = \exp_2(\texttt{gk\_last})$ in code + +Cross-rank state is computed by chaining $\mathbf{M}$ matrices: + +$$\mathbf{S}_r = \mathbf{M}_{r-1} \left(\mathbf{M}_{r-2} \left(\cdots \mathbf{S}_{\text{ext},0} + \mathbf{S}_{\text{ext},1}\right) + \cdots\right) + \mathbf{S}_{\text{ext},r-1}$$ + +**Precision note**: The $\mathbf{M}$ chain multiply must stay in fp32 to avoid accumulated precision loss. In bf16, repeatedly casting fp32 accumulators back to bf16 between iterations causes significant error growth over many chunks. + +--- + +## compress_h0 / expand_h0 + +Optimization for CP mode. Since only the first sequence in the local batch can be a continuation from a previous rank, only its initial state $\mathbf{S}_0 \in \mathbb{R}^{H \times d_k \times d_v}$ is non-zero. +`compress_h0` extracts just that one state to save memory during `save_for_backward`. +`expand_h0` restores the full `[N, H, d_k, d_v]` tensor in backward. diff --git a/fla/ops/cp/README.md b/fla/ops/cp/README.md new file mode 100644 index 0000000000000000000000000000000000000000..bc829d17acbb4fb41b60b1978e8ae26ab2766782 --- /dev/null +++ b/fla/ops/cp/README.md @@ -0,0 +1,69 @@ +# CP Related Features was first implemented by Duyue MA +# And integrated by Zhiyuan Li + +# Context Parallel (CP) Usage Guide + +## Conventions +- CP splits the sequence dimension across ranks. Each rank owns a local chunk + of tokens and runs the operator on that local chunk. +- CP context stores **rank-local** varlen metadata: + - `FLACPContext.cu_seqlens` is rank-local, on GPU (int64 / torch.long). + - `FLACPContext.cu_seqlens_cpu` is rank-local, on CPU (int64 / torch.long). +- Variable-length inputs are represented by global `cu_seqlens` **before** + partition; `build_cp_context` converts them into rank-local metadata. +- CP runs do **not** support `initial_state` or `output_final_state=True`. + +## Build CP Context +```python +from fla.ops.cp import build_cp_context + +# global cu_seqlens before partition (device can be CPU or GPU) +cu_seqlens_global = torch.tensor([0, s1, s1+s2, ..., total], dtype=torch.long, device=device) + +# conv1d_kernel_size is required for causal_conv1d CP path +cp_context = build_cp_context( + cu_seqlens_global, + group=dist.group.WORLD, + conv1d_kernel_size=W, +) +``` + +## Causal Conv1d (CP) +```python +from fla.modules.convolution import causal_conv1d + +# x_local is the rank-local chunk: [1, T_local, D] +y_local, _ = causal_conv1d( + x=x_local, + weight=weight_local, + bias=bias_local, + activation="swish", + cp_context=cp_context, +) +``` +Notes: +- `cp_context` is required. +- `cp_context.conv1d_kernel_size` and `cp_context.cu_seqlens` must be set. +- Do not pass `cu_seqlens`/`cu_seqlens_cpu` manually; they are taken from context. + +## KDA (CP) +```python +from fla.ops.kda import chunk_kda + +o_local, _ = chunk_kda( + q=F.normalize(q_local, p=2, dim=-1), + k=F.normalize(k_local, p=2, dim=-1), + v=v_local, + g=g_local, + beta=beta_local, + cp_context=cp_context, + disable_recompute=disable_recompute, +) +``` +Notes: +- CP expects `B == 1` for varlen and uses rank-local `cu_seqlens` from context. +- `initial_state` and `output_final_state=True` are not supported in CP. + +## Test References +- `tests/context_parallel/test_cp_conv.py` +- `tests/context_parallel/test_cp_kda.py` diff --git a/fla/ops/cp/__init__.py b/fla/ops/cp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0d14e16fd853996234aa0ded83759a489c204068 --- /dev/null +++ b/fla/ops/cp/__init__.py @@ -0,0 +1,25 @@ +# Context Parallel operators and utilities + +from .comm import ( + all_gather_into_tensor, + all_reduce_sum, + conv_cp_send_recv_bwd, + conv_cp_send_recv_fwd, + send_recv_bwd, + send_recv_fwd, +) +from .context import ( + FLACPContext, + build_cp_context, +) + +__all__ = [ + "FLACPContext", + "all_gather_into_tensor", + "all_reduce_sum", + "build_cp_context", + "conv_cp_send_recv_bwd", + "conv_cp_send_recv_fwd", + "send_recv_bwd", + "send_recv_fwd", +] diff --git a/fla/ops/cp/chunk_delta_h.py b/fla/ops/cp/chunk_delta_h.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf76392cd890c064593eb8364ad154e0f3b7cd4 --- /dev/null +++ b/fla/ops/cp/chunk_delta_h.py @@ -0,0 +1,1322 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +import triton +import triton.language as tl + +from fla.ops.cp.comm import all_gather_into_tensor +from fla.ops.utils.op import exp, exp2 +from fla.utils import USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +if TYPE_CHECKING: + from fla.ops.cp.context import FLACPContext + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BV in [32, 64] + ], + key=['H', 'K', 'V', 'BT', 'USE_EXP2'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def pre_process_fwd_kernel_stage1( + k, + v, + w, + g, + gk, + hm, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_h = tl.program_id(0), tl.program_id(1) + i_n = 0 + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + NT = tl.cdiv(T, BT) + + # calculate offset + hm += i_h * K * (K + V) + v += ((bos * H + i_h) * V).to(tl.int64) + k += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + stride_v = H*V + stride_k = H*K + + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + # main recurrence + for i_t in range(NT): + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h1 *= exp2(b_gk_last1)[:, None] + else: + b_h1 *= exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h2 *= exp2(b_gk_last2)[:, None] + else: + b_h2 *= exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h3 *= exp2(b_gk_last3)[:, None] + else: + b_h3 *= exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h4 *= exp2(b_gk_last4)[:, None] + else: + b_h4 *= exp(b_gk_last4)[:, None] + + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v) + + p_h1 = tl.make_block_ptr(hm, (K, V), (K+V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr(hm, (K, V), (K+V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr(hm, (K, V), (K+V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr(hm, (K, V), (K+V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK2': BK2}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BK2 in [32] + ], + key=['H', 'BT', 'FORWARD'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def pre_process_fwd_bwd_kernel_stage2( + k, + w, + g, + gk, + hm, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, + BK1: tl.constexpr, + BK2: tl.constexpr, + FORWARD: tl.constexpr = True, +): + i_k_col, i_h = tl.program_id(0), tl.program_id(1) + i_n = 0 + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + + # calculate offset + hm += i_h * K * (K + V) + k += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + stride_k = H*K + + row = tl.arange(0, BK1) + col = tl.arange(0, BK2) + i_k_col * BK2 + + b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) + for _i_t in range(NT): + if FORWARD: + i_t = _i_t + else: + i_t = NT - 1 - _i_t + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) + elif USE_GK: + b_gk_last = tl.load(gk + (bos + last_idx) * H*K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + if USE_EXP2: + b_gk_last = exp2(b_gk_last) + else: + b_gk_last = exp(b_gk_last) + b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) + else: + b_diag = tl.where(row[:, None] == row[None, :], 1., 0.0) + if FORWARD: + b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) + else: + b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) + b_m_i = b_diag - b_kw + b_m = tl.dot(b_m_i.to(b_w.dtype), b_m.to(b_w.dtype)) + p_m = tl.make_block_ptr(hm + V, (K, K), (K+V, 1), (0, i_k_col * BK2), (BK1, BK2), (1, 0)) + tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def pre_process_fwd_kernel_merged( + k, + v, + w, + g, + gk, + hm, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BK1: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, + MULTI_SEQS: tl.constexpr, +): + i_col, i_h = tl.program_id(0), tl.program_id(1) + if MULTI_SEQS: + i_n = tl.program_id(2) + # Offset hm for this subseq: hm[i_n, h, k, v+k] + hm += i_n * H * K * (K + V) + i_h * K * (K + V) + else: + i_n = 0 + hm += i_h * K * (K + V) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + NT = tl.cdiv(T, BT) + + # Determine if this block handles h (V part) or m (K part) + # i_col is in range [0, cdiv(V + K, BLOCK_SIZE)) + # Columns [0, V) are for h, columns [V, V+K) are for m + is_h_part = i_col * BLOCK_SIZE < V + k += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + stride_k = H * K + + if is_h_part: + # ====== Stage 1: Compute h (K x V) ====== + v += ((bos * H + i_h) * V).to(tl.int64) + stride_v = H * V + i_v = i_col + + # Initialize h accumulators + b_h1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + + # Main recurrence for h + for i_t in range(NT): + # Compute decayed v + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_decay = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_decay += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_decay += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_decay += tl.dot(b_w, b_h4.to(b_w.dtype)) + + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_decay + + last_idx = min((i_t + 1) * BT, T) - 1 + + # Apply g decay + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + # Apply gk decay + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h1 *= exp2(b_gk_last1)[:, None] + else: + b_h1 *= exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h2 *= exp2(b_gk_last2)[:, None] + else: + b_h2 *= exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h3 *= exp2(b_gk_last3)[:, None] + else: + b_h3 *= exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + if USE_EXP2: + b_h4 *= exp2(b_gk_last4)[:, None] + else: + b_h4 *= exp(b_gk_last4)[:, None] + + b_v = b_v.to(k.dtype.element_ty) + + # Update h: h += k^T @ v + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v) + + # Store h results + stride_hm_kv = K + V + p_h1 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + else: + # ====== Stage 2: Compute m (K x K) ====== + # i_col is for m part, map to K dimension + # m starts at column V, so offset = i_col * BLOCK_SIZE - V + # Use tl.cdiv to correctly compute the number of blocks for V dimension + i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) + + # Following stage2 kernel design: + # - BK1 is the full K dimension (next_power_of_2(K)) + # - BLOCK_SIZE is the column block size (like BK2=32 in stage2) + # Each block computes a (BK1, BLOCK_SIZE) sub-matrix of m + row = tl.arange(0, BK1) + col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE + + # Initialize as identity matrix: M_0 = I + b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) + + for i_t in range(NT): + # Load k and w with full BK1 rows + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) + elif USE_GK: + b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + if USE_EXP2: + b_gk_last = exp2(b_gk_last) + else: + b_gk_last = exp(b_gk_last) + b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) + else: + b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) + + # Compute m update: m = (diag - k^T @ w) @ m + b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) + b_m_i = b_diag - b_kw + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) + + # Store m result + stride_hm_kv = K + V + p_m = tl.make_block_ptr(hm + V, (K, K), (stride_hm_kv, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) + tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'HAS_H0': lambda args: args['h0'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BV in [32, 64] + ], + key=['H', 'K', 'V', 'BV'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['pre_or_post_num_ranks', 'rank', 'NUM_SEQ_ENTRIES']) +def merge_fwd_bwd_kernel( + h, # [H, K, V] or [num_non_first, H, K, V] for intracard (or [V, K] when transposed) + ag_hm, # [H, K, K+V] or [S_split, H, K, K+V] for intracard (always [K, V+K]) + pre_or_post_num_ranks, # num_ranks for CP, NUM_SPLIT_SEQS for intracard + rank, # rank for CP, not used for intracard + seq_offsets, # None for CP, [num_split_seqs+1] for intracard + init_offsets, # None for CP, [num_split_seqs+1] for intracard + h0_seq_ids, # None for CP, [num_split_seqs] for intracard + h0, # None or [N_orig, H, K, V] for intracard (or [V, K] when transposed) + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr, + BK: tl.constexpr, + FORWARD: tl.constexpr, # True for FWD, False for BWD + INTRACARD_MODE: tl.constexpr, # True: intracard mode, False: CP mode + NUM_SEQ_ENTRIES, # num_split_seqs for intracard + HAS_H0: tl.constexpr, # Heuristic: whether h0 is provided + TRANSPOSE_STATE: tl.constexpr = False, # When True, h0/h use [V, K] layout; ag_hm always [K, V+K] +): + """ + Unified merge kernel for both CP and Intra-card modes. + + CP mode (INTRACARD_MODE=False): + Grid: (V/BV, H) + Merges across ranks for context parallel. + + Intra-card mode (INTRACARD_MODE=True): + Grid: (V/BV, NUM_SEQ_ENTRIES, H) + Merges across subseqs within card for intra-card context parallel. + + When TRANSPOSE_STATE=True, h0 and output h use [V, K] layout. + ag_hm always uses [K, V+K] layout (from pre_scan). + The recurrence h' = M @ h + he becomes h_T' = h_T @ M^T + he^T. + """ + i_v = tl.program_id(0) + if INTRACARD_MODE: + i_seq = tl.program_id(1) + i_h = tl.program_id(2) + + if i_seq >= NUM_SEQ_ENTRIES: + return + + # Load offsets for this sequence + ss_start = tl.load(seq_offsets + i_seq).to(tl.int32) + ss_end = tl.load(seq_offsets + i_seq + 1).to(tl.int32) + init_base = tl.load(init_offsets + i_seq).to(tl.int32) + num_subseqs = ss_end - ss_start + + stride_hm_s = H * K * (V + K) + stride_hm_h = K * (V + K) + + # Initialize from h0 if provided + if HAS_H0: + orig_seq_id = tl.load(h0_seq_ids + i_seq).to(tl.int32) + if TRANSPOSE_STATE: + p_h0 = tl.make_block_ptr( + h0 + (orig_seq_id * H + i_h) * V * K, + (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) + ) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + else: + p_h0 = tl.make_block_ptr( + h0 + (orig_seq_id * H + i_h) * K * V, + (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) + ) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + else: + if TRANSPOSE_STATE: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + # Merge loop over subseqs + for idx in range(num_subseqs): + i_ss = ss_start + idx + base = i_ss * stride_hm_s + i_h * stride_hm_h + + # he and m are always in [K, V+K] layout from pre_scan + p_he = tl.make_block_ptr( + ag_hm + base, (K, V), (V + K, 1), (0, i_v * BV), (BK, BV), (1, 0) + ) + b_he = tl.load(p_he, boundary_check=(0, 1)).to(tl.float32) + p_m = tl.make_block_ptr( + ag_hm + base + V, (K, K), (V + K, 1), (0, 0), (BK, BK), (1, 0) + ) + b_m = tl.load(p_m, boundary_check=(0, 1)).to(tl.float32) + if TRANSPOSE_STATE: + # h_T' = h_T @ M^T + he^T + b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_m)) + tl.trans(b_he) + else: + b_h = tl.dot(b_m.to(tl.float32), b_h.to(tl.float32)) + b_he.to(tl.float32) + + # Store for non-first subseqs + if idx < num_subseqs - 1: + init_idx = init_base + idx + stride_init = H * K * V + if TRANSPOSE_STATE: + p_out = tl.make_block_ptr( + h + init_idx * stride_init + i_h * V * K, + (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) + ) + else: + p_out = tl.make_block_ptr( + h + init_idx * stride_init + i_h * K * V, + (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) + ) + tl.store(p_out, b_h.to(p_out.dtype.element_ty), boundary_check=(0, 1)) + else: + # CP mode + i_h = tl.program_id(1) + num_ranks = pre_or_post_num_ranks.to(tl.int32) + h += i_h * K * V + ag_hm += i_h * K * (K + V) + stride = H * K * (K + V) + if TRANSPOSE_STATE: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + for idx in range(num_ranks): + if FORWARD: + cur_rank = rank - num_ranks + idx + else: + cur_rank = rank + num_ranks - idx + p_ag_h = tl.make_block_ptr(ag_hm + cur_rank * stride, (K, V), (K + V, 1), (0, i_v * BV), (BK, BV), (1, 0)) + b_ag_h = tl.load(p_ag_h, boundary_check=(0, 1)) + p_ag_m = tl.make_block_ptr(ag_hm + cur_rank * stride + V, (K, K), (K + V, 1), (0, 0), (BK, BK), (1, 0)) + b_ag_m = tl.load(p_ag_m, boundary_check=(0, 1)) + if TRANSPOSE_STATE: + b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_ag_m).to(tl.float32)) + tl.trans(b_ag_h).to(tl.float32) + else: + b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32)) + b_ag_h.to(tl.float32) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) + for BV in [64, 32] + ], + key=['H', 'K', 'V', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def pre_process_bwd_kernel_stage1( + q, + k, + w, + g, + gk, + do, + dhm, + dv, + cu_seqlens, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_h = tl.program_id(0), tl.program_id(1) + i_n = 0 + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + NT = tl.cdiv(T, BT) + + # [BK, BV] + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += ((bos * H + i_h) * K).to(tl.int64) + k += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + do += ((bos * H + i_h) * V).to(tl.int64) + dv += ((bos * H + i_h) * V).to(tl.int64) + dhm += i_h * K * (V + K) + + stride_v = H*V + stride_k = H*K + + for i_t in range(NT - 1, -1, -1): + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) + bg_last_exp = exp(bg_last) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g_exp = exp(b_g) + + p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + # Update dv + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + + if K > 128: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + + if K > 192: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, boundary_check=(0, 1)) + + # Update dh + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + b_dh1 *= exp(b_gk_last1[:, None]) + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 64: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + b_dh2 *= exp(b_gk_last2[:, None]) + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 128: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + b_dh3 *= exp(b_gk_last3[:, None]) + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 192: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + b_dh4 *= exp(b_gk_last4[:, None]) + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_dh2 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_dh3 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_dh4 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) + ], + key=['H', 'K', 'V', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def pre_process_bwd_kernel_merged( + q, + k, + w, + g, + gk, + do, + dhm, + dv, + cu_seqlens, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BK1: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Merged backward kernel that computes both dh (K x V) and dm (K x K) in a single kernel. + + Similar to pre_process_fwd_kernel_merged, this kernel uses a unified grid where: + - Columns [0, V) are for computing dh (stage 1) + - Columns [V, V+K) are for computing dm (stage 2) + """ + i_col, i_h = tl.program_id(0), tl.program_id(1) + i_n = 0 + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + NT = tl.cdiv(T, BT) + + # Determine if this block handles dh (V part) or dm (K part) + is_dh_part = i_col * BLOCK_SIZE < V + + # Calculate offsets + q += ((bos * H + i_h) * K).to(tl.int64) + k += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + dhm += i_h * K * (V + K) + stride_k = H * K + + if is_dh_part: + # ====== Stage 1: Compute dh (K x V) ====== + do += ((bos * H + i_h) * V).to(tl.int64) + dv += ((bos * H + i_h) * V).to(tl.int64) + stride_v = H * V + i_v = i_col + + # Initialize dh accumulators + b_dh1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + + # Main recurrence for dh (reverse order) + for i_t in range(NT - 1, -1, -1): + last_idx = min((i_t + 1) * BT, T) - 1 + + if USE_G: + # Note: pre_process_bwd_kernel_stage1 always uses exp for USE_G, + # regardless of USE_EXP2. This is for consistency with the original design. + bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) + bg_last_exp = exp(bg_last) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g_exp = exp(b_g) + + p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + # Update dv + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k1 = tl.arange(0, 64) + if USE_EXP2: + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + else: + b_gk_last1 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k2 = 64 + o_k1 + if USE_EXP2: + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + else: + b_gk_last2 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + + if K > 128: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k3 = 128 + o_k1 + if USE_EXP2: + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + else: + b_gk_last3 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + + if K > 192: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_GK: + o_k4 = 192 + o_k1 + if USE_EXP2: + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + else: + b_gk_last4 = tl.load(gk + (bos + last_idx) * H * K + i_h * K + + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) + b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + # Note: pre_process_bwd_kernel_stage1 always uses exp for USE_G + b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, boundary_check=(0, 1)) + + # Update dh + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if USE_EXP2: + b_dh1 *= exp2(b_gk_last1[:, None]) + else: + b_dh1 *= exp(b_gk_last1[:, None]) + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + if K > 64: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if USE_EXP2: + b_dh2 *= exp2(b_gk_last2[:, None]) + else: + b_dh2 *= exp(b_gk_last2[:, None]) + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + if K > 128: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if USE_EXP2: + b_dh3 *= exp2(b_gk_last3[:, None]) + else: + b_dh3 *= exp(b_gk_last3[:, None]) + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + if K > 192: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if USE_EXP2: + b_dh4 *= exp2(b_gk_last4[:, None]) + else: + b_dh4 *= exp(b_gk_last4[:, None]) + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) + + # Store dh results + p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_dh2 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_dh3 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_dh4 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + else: + # ====== Stage 2: Compute dm (K x K) ====== + # i_col is for dm part, map to K dimension + i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) + + # Following stage2 kernel design for backward (FORWARD=False) + # - BK1 is the full K dimension (next_power_of_2(K)) + # - BLOCK_SIZE is the column block size + row = tl.arange(0, BK1) + col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE + + # Initialize as identity matrix: M_0 = I + b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) + + for _i_t in range(NT): + # Reverse order for backward + i_t = NT - 1 - _i_t + + # Load k and w with full BK1 rows + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if USE_EXP2: + b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + else: + b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) + elif USE_GK: + b_gk_last = tl.load(gk + (bos + last_idx) * H * K + i_h * K + row, mask=(row < K), other=0.).to(tl.float32) + if USE_EXP2: + b_gk_last = exp2(b_gk_last) + else: + b_gk_last = exp(b_gk_last) + b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) + else: + b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) + + # Compute dm update for backward: m = (diag - w^T @ k) @ m + # Note: FORWARD=False uses tl.trans(b_w) @ b_k instead of tl.trans(b_k) @ b_w + b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) + b_m_i = b_diag - b_kw + # Keep m chain in fp32 to avoid precision loss from repeated bf16 casting + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) + + # Store dm result + p_m = tl.make_block_ptr(dhm + V, (K, K), (V + K, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) + tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_h_pre_process( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, + initial_state: torch.Tensor | None = None, + context: FLACPContext = None, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + if context is None or context.group is None: + return initial_state + assert initial_state is None, "When enable CP, the provided initial_state must be None." + rank = dist.get_rank(group=context.group) + + B, T, H, K, V = *k.shape, u.shape[-1] + BT = chunk_size + BK = triton.next_power_of_2(K) + + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N = B + else: + N = len(cu_seqlens) - 1 + assert K <= 256, "current kernel does not support head dimension larger than 256." + + hm = k.new_zeros(H, K, (V + K), dtype=torch.float32) + if transpose_state_layout: + initial_state = k.new_zeros(N, H, V, K, dtype=torch.float32) + else: + initial_state = k.new_zeros(N, H, K, V, dtype=torch.float32) + if not context.is_last_rank: + BLOCK_SIZE = 32 if K <= 64 else 64 + grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) + pre_process_fwd_kernel_merged[grid]( + k=k, + v=u, + w=w, + g=g, + gk=gk, + hm=hm, + cu_seqlens=cu_seqlens[-2:], + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK1=BK, + USE_EXP2=use_exp2, + BLOCK_SIZE=BLOCK_SIZE, + MULTI_SEQS=False, + ) + ag_hm, _ = all_gather_into_tensor(hm, group=context.group) + if not context.is_first_rank: + def grid(meta): return (triton.cdiv(V, meta['BV']), H) + merge_fwd_bwd_kernel[grid]( + h=initial_state[0], + ag_hm=ag_hm, + pre_or_post_num_ranks=context.pre_num_ranks, + rank=rank, + seq_offsets=None, + init_offsets=None, + h0_seq_ids=None, + h0=None, + H=H, + K=K, + V=V, + BK=BK, + FORWARD=True, + INTRACARD_MODE=False, + NUM_SEQ_ENTRIES=0, + TRANSPOSE_STATE=transpose_state_layout, + ) + return initial_state + + +def chunk_gated_delta_rule_bwd_dhu_pre_process( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, + dht: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + context: FLACPContext | None = None, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if context is None or context.group is None: + return dht, initial_state + assert dht is None, "When enable CP, the provided dht must be None." + rank = dist.get_rank(context.group) + + B, T, H, K, V = *q.shape, do.shape[-1] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = 64 + assert K <= 256, "current kernel does not support head dimension being larger than 256." + BK = triton.next_power_of_2(K) + + if cu_seqlens is None: + N = B + else: + N = len(cu_seqlens) - 1 + + dhm = q.new_zeros(H, K, V + K, dtype=torch.float32) + if transpose_state_layout: + dht = q.new_zeros(N, H, V, K, dtype=torch.float32) + else: + dht = q.new_zeros(N, H, K, V, dtype=torch.float32) + + if not context.is_first_rank: + BLOCK_SIZE = 32 if K <= 64 else 64 + grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) + pre_process_bwd_kernel_merged[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + do=do, + dhm=dhm, + dv=dv, + cu_seqlens=cu_seqlens[:2], + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK1=BK, + USE_EXP2=use_exp2, + BLOCK_SIZE=BLOCK_SIZE, + ) + + ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) + + if not context.is_last_rank: + def grid(meta): return (triton.cdiv(V, meta['BV']), H) + merge_fwd_bwd_kernel[grid]( + h=dht[-1], + ag_hm=ag_dhm, + pre_or_post_num_ranks=context.post_num_ranks, + rank=rank, + seq_offsets=None, + init_offsets=None, + h0_seq_ids=None, + h0=None, + H=H, + K=K, + V=V, + BK=BK, + FORWARD=False, + INTRACARD_MODE=False, + NUM_SEQ_ENTRIES=0, + TRANSPOSE_STATE=transpose_state_layout, + ) + + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + return dht, None + + +def compress_h0(h0: torch.Tensor, context: FLACPContext): + if h0 is None or len(context.cu_seqlens) == 2: + return h0 + # Here must use clone op or the full tensor will be saved for backward + return h0[:1].clone() + + +def expand_h0(h0: torch.Tensor, context: FLACPContext): + if h0 is None or len(context.cu_seqlens) == 2: + return h0 + B = len(context.cu_seqlens) - 1 + expand_h0 = h0.new_zeros(B, *h0.shape[1:]) + expand_h0[:1] = h0 + return expand_h0 diff --git a/fla/ops/cp/comm.py b/fla/ops/cp/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..7b6cc339554c5a7948a4917e805d6165d475c0fc --- /dev/null +++ b/fla/ops/cp/comm.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + +def all_gather_into_tensor( + inp: torch.Tensor, + out: torch.Tensor | None = None, + group: ProcessGroup | None = None, + async_op: bool = False +) -> tuple[torch.Tensor, dist.Work | None]: + """ + All-gather a tensor across ranks. + + Args: + inp: Input tensor to gather + out: Optional output tensor of shape [world_size, *inp.shape] + group: Process group + async_op: Whether to perform async operation + + Returns: + Tuple of (output tensor, handle if async_op else None) + """ + world_size = dist.get_world_size(group=group) + if out is None: + out = torch.empty(world_size, *inp.shape, device=inp.device, dtype=inp.dtype) + handle = dist.all_gather_into_tensor(out, inp, group=group, async_op=async_op) + return out, handle + + +def all_reduce_sum( + inp: torch.Tensor, + group: ProcessGroup | None = None, + async_op: bool = False +) -> tuple[torch.Tensor, dist.Work | None]: + """ + All-reduce sum a tensor across ranks. + + Args: + inp: Input tensor to reduce (modified in-place) + group: Process group + async_op: Whether to perform async operation + + Returns: + Tuple of (reduced tensor, handle if async_op else None) + """ + handle = dist.all_reduce(inp, op=dist.ReduceOp.SUM, group=group, async_op=async_op) + return inp, handle + + +def send_recv_fwd( + send_tensor: torch.Tensor, + group: ProcessGroup, + recv_from_prev: bool = True +) -> torch.Tensor: + """ + Forward pass communication: send tensor to next rank, receive from previous rank. + + Uses all_gather for simplicity and to ensure all ranks participate. + + Args: + send_tensor: Tensor to send (e.g., tails for conv1d) + group: Process group + recv_from_prev: If True, receive from previous rank; if False, receive from next rank + + Returns: + Received tensor from the specified rank (zeros if no valid source) + """ + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + + # All-gather to ensure all ranks participate + gathered, _ = all_gather_into_tensor(send_tensor, group=group, async_op=False) + + if recv_from_prev: + # Receive from previous rank + if rank == 0: + return torch.zeros_like(send_tensor) + else: + return gathered[rank - 1].clone() + else: + # Receive from next rank + if rank == world_size - 1: + return torch.zeros_like(send_tensor) + else: + return gathered[rank + 1].clone() + + +def send_recv_bwd( + send_tensor: torch.Tensor, + group: ProcessGroup, + recv_from_next: bool = True +) -> torch.Tensor: + """ + Backward pass communication: send gradient to previous rank, receive from next rank. + + Uses all_gather for simplicity and to ensure all ranks participate. + + Args: + send_tensor: Gradient tensor to send + group: Process group + recv_from_next: If True, receive from next rank; if False, receive from previous rank + + Returns: + Received gradient tensor from the specified rank (zeros if no valid source) + """ + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + + # All-gather to ensure all ranks participate + gathered, _ = all_gather_into_tensor(send_tensor, group=group, async_op=False) + + if recv_from_next: + # Receive from next rank + if rank == world_size - 1: + return torch.zeros_like(send_tensor) + else: + return gathered[rank + 1].clone() + else: + # Receive from previous rank + if rank == 0: + return torch.zeros_like(send_tensor) + else: + return gathered[rank - 1].clone() + + +# ============ Convenience aliases for conv1d CP ============ + +def conv_cp_send_recv_fwd(tails: torch.Tensor, group: ProcessGroup) -> torch.Tensor: + """ + Conv1d CP forward: each rank sends its tails, receives previous rank's tails as heads. + + Args: + tails: [W-1, D] or [N, D, W-1] - tail tokens from current rank + group: Process group + + Returns: + heads: Same shape as tails - head tokens from previous rank (zeros for rank 0) + """ + return send_recv_fwd(tails, group, recv_from_prev=True) + + +def conv_cp_send_recv_bwd(d_initial_state: torch.Tensor, group: ProcessGroup) -> torch.Tensor: + """ + Conv1d CP backward: each rank sends d_initial_state, receives from next rank. + + The received gradient should be added to the last W-1 tokens' gradient. + + Args: + d_initial_state: [W-1, D] or [N, D, W-1] - gradient w.r.t. initial state + group: Process group + + Returns: + recv_grad: Same shape - gradient from next rank (zeros for last rank) + """ + return send_recv_bwd(d_initial_state, group, recv_from_next=True) diff --git a/fla/ops/cp/context.py b/fla/ops/cp/context.py new file mode 100644 index 0000000000000000000000000000000000000000..0e9a80e204732983e1610246d801d6b77656aa03 --- /dev/null +++ b/fla/ops/cp/context.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +from fla.utils import tensor_cache + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + +@dataclass +class FLACPContext: + """FLA Context Parallel Context - Operator-level context management.""" + group: ProcessGroup | None = None + cu_seqlens: torch.Tensor | None = None + cu_seqlens_cpu: torch.Tensor | None = None + is_last_rank: bool | None = None + pre_num_ranks: int | None = None + is_first_rank: bool | None = None + post_num_ranks: int | None = None + conv1d_kernel_size: int | None = None + pre_num_conv_tokens: int | None = None + + def copy_for_backward(self) -> FLACPContext: + """Create a copy for backward pass (useful when PP_SIZE > 1).""" + return FLACPContext( + group=self.group, + cu_seqlens=self.cu_seqlens.clone() if self.cu_seqlens is not None else None, + cu_seqlens_cpu=self.cu_seqlens_cpu.clone() if self.cu_seqlens_cpu is not None else None, + is_last_rank=self.is_last_rank, + pre_num_ranks=self.pre_num_ranks, + is_first_rank=self.is_first_rank, + post_num_ranks=self.post_num_ranks, + conv1d_kernel_size=self.conv1d_kernel_size, + pre_num_conv_tokens=self.pre_num_conv_tokens, + ) + + @property + def num_seqs(self) -> int: + """Number of sequences in this rank.""" + return 0 if self.cu_seqlens is None else len(self.cu_seqlens) - 1 + + @property + def is_cp_enabled(self) -> bool: + """Whether context parallel is enabled.""" + return self.group is not None + + +@tensor_cache +def get_cp_cu_seqlens( + cu_seqlens: torch.LongTensor, + cu_seqlens_cpu: torch.LongTensor | None = None, + world_size: int | None = None, + rank: int | None = None, + group: dist.ProcessGroup | None = None, + conv1d_kernel_size: int | None = None +) -> FLACPContext: + # 1. Initialize environment info + if world_size is None: + assert group is not None + world_size = dist.get_world_size(group=group) + rank = dist.get_rank(group=group) + + # 2. Operate on CPU to avoid D2H sync and leverage vectorization (int64/long) + if cu_seqlens_cpu is None: + cu_seqlens_cpu = cu_seqlens.cpu() + cu_seqlens_cpu = cu_seqlens_cpu.to(dtype=torch.long) + + # Get total tokens and current rank's responsible range + # Assume cu_seqlens is [0, s1, s1+s2, ..., total] + total_tokens = cu_seqlens_cpu[-1].item() + part_len = total_tokens // world_size + rank_start = part_len * rank + rank_end = rank_start + part_len + + # 3. Vectorized search: find sequences overlapping with current rank's interval [rank_start, rank_end) + # We need to find idx such that: global_ends[idx] > rank_start AND global_starts[idx] < rank_end + + # Optimization: cu_seqlens is sorted, use searchsorted to quickly locate boundaries + # Find first sequence whose end > rank_start + # cu_seqlens_cpu[1:] contains all sequence end points + start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], rank_start, side='right') + + # Find first sequence whose start >= rank_end, sequences before this may overlap + # cu_seqlens_cpu[:-1] contains all sequence start points + end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], rank_end, side='left') + + # Slice cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] to get relevant global cu_seqlens nodes + # +1 because end_seq_idx is an open boundary, and cu_seqlens length is num_seqs + 1 + subset_cu_seqlens = cu_seqlens_cpu[start_seq_idx: end_seq_idx + 1] + + # 4. Compute local cu_seqlens on CPU (int32) + # Clamp global coordinates to [rank_start, rank_end], subtract rank_start to get local coordinates + # unique_consecutive removes duplicates from clamping (e.g., sequences entirely outside this rank) + local_cu_seqlens_cpu = ( + subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start + ).unique_consecutive().to(torch.int32) + + # Transfer to GPU (int32, small tensor, fast transfer) + # non_blocking=True can further hide latency in CUDA streams + local_cu_seqlens_gpu = local_cu_seqlens_cpu.to( + device=cu_seqlens.device, non_blocking=True + ) + + # 5. Compute Context Parallel metadata (first/last rank info) + # Use slice endpoints directly, avoiding loops + + # Get global info for the first sequence that has data on current rank + first_seq_global_start = cu_seqlens_cpu[start_seq_idx].item() + # Get global info for the last sequence that has data on current rank + last_seq_global_end = cu_seqlens_cpu[end_seq_idx].item() + + # Number of tokens current rank needs from previous ranks for conv + pre_num_conv_tokens = max(0, rank_start - first_seq_global_start) + + # Compute first sequence's starting rank + first_rank_of_first_seq = first_seq_global_start // part_len + # Number of previous ranks current rank needs to receive state from + pre_num_ranks = rank - first_rank_of_first_seq + # Whether current rank is the first in the sequence's processing chain + is_first_rank = (rank == first_rank_of_first_seq) + + # Compute last sequence's ending rank + # (last_seq_global_end - 1) is the index of the last token + last_rank_of_last_seq = (last_seq_global_end - 1) // part_len + # Number of subsequent ranks current rank needs to send state to + post_num_ranks = last_rank_of_last_seq - rank + # Whether current rank is the last in the sequence's processing chain + is_last_rank = (rank == last_rank_of_last_seq) + + return FLACPContext( + group=group, + cu_seqlens=local_cu_seqlens_gpu, + cu_seqlens_cpu=local_cu_seqlens_cpu, + is_last_rank=is_last_rank, + pre_num_ranks=pre_num_ranks, + is_first_rank=is_first_rank, + post_num_ranks=post_num_ranks, + conv1d_kernel_size=conv1d_kernel_size, + pre_num_conv_tokens=pre_num_conv_tokens + ) + + +def build_cp_context( + cu_seqlens: torch.Tensor, + group: ProcessGroup, + conv1d_kernel_size: int | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, +) -> FLACPContext: + """Build a CP context for the given cu_seqlens and process group. + + Args: + cu_seqlens: Cumulative sequence lengths tensor (before partition). + group: Process group for CP communication. + conv1d_kernel_size: Kernel size for convolution (optional). + cu_seqlens_cpu: CPU version of cu_seqlens to avoid d2h transfer (optional). + + Returns: + FLACPContext with computed cu_seqlens and rank information. + """ + return get_cp_cu_seqlens(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, group=group, conv1d_kernel_size=conv1d_kernel_size) diff --git a/fla/ops/delta_rule/README.md b/fla/ops/delta_rule/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7dfe2f01bd140b9464c072c4732808bd38d471e6 --- /dev/null +++ b/fla/ops/delta_rule/README.md @@ -0,0 +1,90 @@ +# Chunkwise-form Parallelism of DeltaNet + +This section expands on the formulation presented in Appendix B of the DeltaNet paper.[^1] + +To reduce notational clutter, we focus on the first chunk, denoting $\mathbf{S}^r=\mathbf{S}_{[1]}^r$. By partially expanding the recurrence, we have: +```math +\begin{equation} +\begin{aligned} +\mathbf{S}^r &= \underbrace{\left(\prod_{i=1}^r \mathbf{I} - \beta^i \bf{k}^i \bf{k}^{i\top} \right)}_{:= \mathbf{P}^r} \cdot\mathbf{S}^{0} + \overbrace{\sum_{i=1}^{r} \underbrace{\left(\prod_{j=i+1}^r \mathbf{I} - \beta^j \bf{k}^j \bf{k}^{j\top} \right)}_{:= \mathbf{P}_{i+1}^r}\beta^i \bf{k}^i\bf{v}^{i\top}}^{:=\mathbf{H}^r} \\ +&=\mathbf{P}^r \cdot \mathbf{S}^{0} + \mathbf{H}^r +\end{aligned} +\end{equation} +``` + +where $\mathbf{P}_i^r$ involves cumulative products of generalized Householder matrices. +We abbreviate $\mathbf{P}_1^r$ as $\mathbf{P}^r$. +This can be optimized using the classical WY representation: +```math +\begin{equation} +\mathbf{P}^{r} = \mathbf{I} - \sum_{i=1}^{r}\bf{k}^i\bf{w}^{i\top} \in \mathbb{R}^{d_k \times d_k};\qquad +\bf{w}^r = \beta^r \left(\bf{k}^r - \sum_{i=1}^{r-1} \left(\bf{k}^{r\top}\bf{k}^i \right)\bf{w}^i \right) \in \mathbb{R}^{d_k} +\end{equation} +``` + +We prove this by induction: +```math +\begin{align*} +\mathbf{P}^{r} &= \prod_{i=1}^r \mathbf{I} - \beta^i \bf{k}^i \bf{k}^{i\top} \\ +&= \left(\mathbf{I} - \beta^r \bf{k}^r \bf{k}^{r\top}\right)\mathbf{P}^{r-1} \\ +&= \left(\mathbf{I} - \beta^r \bf{k}^r \bf{k}^{r\top}\right)\left(\mathbf{I} - \sum_{i=1}^{r-1}\bf{k}^i\bf{w}^{i\top}\right) \\ +&= \mathbf{I} - \sum_{i=1}^{r-1}\bf{k}^i\bf{w}^{i\top} - \beta^r \bf{k}^r \bf{k}^{r\top} + \beta^r\bf{k}^r \bf{k}^{r\top} \left(\sum_{i=1}^{r-1}\bf{k}^i\bf{w}^{i\top}\right) \\ +&= \mathbf{I} - \sum_{i=1}^{r-1}\bf{k}^i\bf{w}^{i\top} - \beta^r \bf{k}^r \left(\bf{k}^{r} - \left(\sum_{i=1}^{r-1}\left(\bf{k}^{r\top} \bf{k}^i\right)\bf{w}^{i}\right) \right)^\top \\ +&= \mathbf{I} - \sum_{i=1}^{r}\bf{k}^i\bf{w}^{i\top} +\end{align*} +``` + +Similarly, $\mathbf{H}^r$ can be represented as: +```math +\begin{equation} +\mathbf{H}^{r} = \sum_{i=1}^{r} \bf{k}^i \bf{u}^{i\top} \in \mathbb{R}^{d_k \times d_v};\qquad \bf{u}^r = \beta^r \left(\bf{v}^r - \sum_{i=1}^{r-1} \left(\bf{k}^{r\top}\bf{k}^i\right) \bf{u}^i \right)\in \mathbb{R}^{d_v} +\end{equation} +``` + +This can also be proven by induction: +```math +\begin{align*} +\mathbf{H}^{r} &= \sum_{i=1}^{r} \mathbf{P}_{i+1}^r \beta^i \bf{k}^i \bf{v}^{i\top}\\ +&= \left(\mathbf{I} - \beta^r \bf{k}^r \bf{k}^{r\top}\right) \mathbf{H}^{r-1} + \beta^r \bf{k}^r \bf{v}^{r\top}\\ +&= \sum_{i=1}^{r-1}\bf{k}^i \bf{u}^{i\top} - \beta^r \bf{k}^r \bf{k}^{r\top} \sum_{i=1}^{r-1}\bf{k}^i \bf{u}^{i\top} +\beta^r \bf{k}^r \bf{v}^{r\top}\\ +&= \sum_{i=1}^{r-1}\bf{k}^i \bf{u}^{i\top} + \bf{k}^r \left(\beta^r \bf{v}^{r\top}-\beta^r \bf{k}^{r\top} \sum_{i=1}^{r-1}\bf{k}^i \bf{u}^{i\top}\right) \\ +&= \sum_{i=1}^{r-1}\bf{k}^i \bf{u}^{i\top} + \bf{k}^r \beta^r\left(\bf{v}^{r}-\sum_{i=1}^{r-1}\left(\bf{k}^{r\top}\bf{k}^{i}\right)\bf{u}^{i} \right)^\top \\ +&=\sum_{i=1}^{r} \bf{k}^i \bf{u}^{i\top} +\end{align*} +``` + +In matrix form, $\mathbf{P}$ and $\mathbf{H}$ can be written as: +```math +\begin{equation} +\mathbf{P}=\mathbf{I}-\mathbf{K}^\top\mathbf{W} \in \mathbb{R}^{d_k \times d_k}, \qquad\mathbf{H}=\mathbf{K}^\top\mathbf{U} \in \mathbb{R}^{d_k\times d_v} +\end{equation} +``` + +Now we can derive the matrix form of $\mathbf{W}$ and $\mathbf{U}$: +```math +\begin{align*} +\mathbf{W} &= \mathrm{diag}(\beta) \mathbf{K} - \mathrm{tril}(\mathrm{diag}(\beta) \mathbf{K}\mathbf{K}^\top, -1)\mathbf{W}\\ +\left(\mathbf{I} + \mathrm{tril}(\mathrm{diag}(\beta) \mathbf{K}\mathbf{K}^\top, -1)\right) \mathbf{W} &= \mathrm{diag}(\beta) \mathbf{K} +\end{align*} +``` +A similar process holds for $\mathbf{U}$. We can further write $\mathbf{W}$ and $\mathbf{U}$ in matrix form: +```math +\begin{align*} +\mathbf{T} &= \left(\mathbf{I} + \mathrm{tril}\left(\mathrm{diag}(\beta)\mathbf{K} \mathbf{K}^\top,-1\right)\right)^{-1}\mathrm{diag}\left(\beta\right)\in \mathbb{R}^{C \times C}\\ +\mathbf{W} &= \mathbf{T} \mathbf{K}\in \mathbb{R}^{C \times d_k}\\ +\mathbf{U} &= \mathbf{T}\mathbf{V}\in \mathbb{R}^{C \times d_v} +\end{align*} +``` + +Substituting these back into the original equations yields a hardware-efficient chunkwise algorithm for DeltaNet that leverages matrix multiplications, enabling tensor core based GPU optimization: +```math +\begin{equation} +\begin{aligned} +\mathbf{S} &= \mathbf{P}\cdot\mathbf{S}^0 + \mathbf{H} \\ +&= \mathbf{S}^0 + \mathbf{K}^\top (\mathbf{U} -\mathbf{W} \mathbf{S}^0) \in \mathbb{R}^{d_k \times d_v}\\ +\mathbf{O} &= \mathbf{Q} \mathbf{S}^0 + (\mathbf{Q} \mathbf{K}^{\top} \odot \mathbf{M}) \left(\mathbf{U} - \mathbf{W} \mathbf{S}^0\right) \in \mathbb{R}^{C \times d_v} +\end{aligned} +\end{equation} +``` + +[^1]: https://arxiv.org/abs/2406.06484 diff --git a/fla/ops/delta_rule/__init__.py b/fla/ops/delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..664658684f8e0e8d78467219979740a4bc5ef6c9 --- /dev/null +++ b/fla/ops/delta_rule/__init__.py @@ -0,0 +1,10 @@ + +from .chunk import chunk_delta_rule +from .fused_chunk import fused_chunk_delta_rule +from .fused_recurrent import fused_recurrent_delta_rule + +__all__ = [ + 'fused_chunk_delta_rule', + 'fused_recurrent_delta_rule', + 'chunk_delta_rule', +] diff --git a/fla/ops/delta_rule/chunk.py b/fla/ops/delta_rule/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..c84430b151fc279a299d6f3d9a6253d33f5bbf54 --- /dev/null +++ b/fla/ops/delta_rule/chunk.py @@ -0,0 +1,327 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.delta_rule.wy_fast import prepare_wy_repr_bwd, prepare_wy_repr_fwd, recompute_w_u_fwd +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + # obtain WY representation. u is actually the new v. + w, u, A = prepare_wy_repr_fwd( + k=k, + v=v, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=None, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=None, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return o, A, final_state + + +def chunk_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=None, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + do=do, + g=None, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=None, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dq, dk, dw, _ = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + h=h, + w=w, + dv=dv, + do=do, + dh=dh, + g=None, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk2, dv, db = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + return dq, dk, dv, db, dh0 + + +class ChunkDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + ): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + else: + q_rstd, k_rstd = None, None + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + o, A, final_state = chunk_delta_rule_fwd( + q=q, + k=k, + v=v, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, beta, A, initial_state, cu_seqlens, chunk_indices) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + q, q_rstd, k, k_rstd, v, beta, A, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors + + dq, dk, dv, db, dh0 = chunk_delta_rule_bwd( + q=q, + k=k, + v=v, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), db.to(beta.dtype), None, dh0, None, None, None, None + + +@torch.compiler.disable +def chunk_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use qk l2norm within the kernel for saving GPU memory. + Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.delta_rule import chunk_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + assert q.dtype == k.dtype == v.dtype + assert q.dtype != torch.float32, "ChunkDeltaRuleFunction does not support float32. Please use bfloat16." + assert len(beta.shape) == 3, "beta must be of shape (batch size, num of head, seq len)." + + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + scale = k.shape[-1] ** -0.5 if scale is None else scale + o, final_state = ChunkDeltaRuleFunction.apply( + q, + k, + v, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/delta_rule/fused_chunk.py b/fla/ops/delta_rule/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..36231fa608d474ce3406da7b8fe089a8343187c9 --- /dev/null +++ b/fla/ops/delta_rule/fused_chunk.py @@ -0,0 +1,5 @@ + +def fused_chunk_delta_rule( + **kwargs, +): + raise NotImplementedError("fused_chunk_delta_rule is deprecated. Please use chunk_delta_rule instead.") diff --git a/fla/ops/delta_rule/fused_recurrent.py b/fla/ops/delta_rule/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..f944f02903efced2f44ec444b4885bc5f2c47a03 --- /dev/null +++ b/fla/ops/delta_rule/fused_recurrent.py @@ -0,0 +1,533 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_delta_rule_fwd_kernel( + q, + k, + v, + u, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_u = u + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + bos * H + i_h + p_o = o + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + + mask_k = (i_k * BK + tl.arange(0, BK)) < K + mask_v = (i_v * BV + tl.arange(0, BV)) < V + mask_h = mask_k[None, :] & mask_v[:, None] + + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_v_minus = tl.sum(b_h * b_k[None, :], axis=1) + b_v -= b_v_minus + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + tl.store(p_u, b_v.to(p_v.dtype.element_ty), mask=mask_v) + b_v *= b_beta + b_h += b_k[None, :] * b_v[:, None] + b_o = b_h * b_q[None, :] + b_o = tl.sum(b_o, axis=1) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_o += H*V + p_v += H*V + p_u += H*V + p_beta += H * (V if IS_BETA_HEADWISE else 1) + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_delta_rule_bwd_kernel( + q, + k, + v, + beta, + h0, + dh0, + dht, + do, + dq, + dk, + dv, + db, + cu_seqlens, + scale, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NK: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar + USE_INITIAL_STATE: tl.constexpr, # whether to use dh0 + USE_FINAL_STATE_GRADIENT: tl.constexpr, # whether to use dht + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + mask_k = i_k * BK + tl.arange(0, BK) < K + mask_v = i_v * BV + tl.arange(0, BV) < V + + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + if IS_BETA_HEADWISE: + p_beta = beta + (bos + T - 1) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + p_dbeta = db + ((i_v * NK + i_k) * all + bos + T - 1) * H*V + i_h * V + tl.arange(0, BV) + else: + p_beta = beta + (bos + T - 1) * H + i_h + p_dbeta = db + (i_v * all + bos + T - 1) * H + i_h + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_ht = dht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + b_dh += tl.load(p_ht, mask=mask_k[:, None] & mask_v[None, :], other=0).to(tl.float32) + + for _ in range(T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_dh += b_q[:, None] * b_do[None, :] + b_dk = tl.sum(b_dh * (b_v * b_beta)[None, :], axis=1) + b_dv = tl.sum(b_dh * b_k[:, None], axis=0) + + b_db = b_dv * b_v if IS_BETA_HEADWISE else tl.sum(b_dv * b_v) + b_dv = b_dv * b_beta + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) + if IS_BETA_HEADWISE: + tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty), mask=mask_v) + else: + tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty)) + + b_dh -= b_k[:, None] * b_dv[None, :] + + p_q -= H*K + p_k -= H*K + p_v -= H*V + p_do -= H*V + p_dk -= H*K + p_dv -= H*V + p_dbeta -= H * (V if IS_BETA_HEADWISE else 1) + p_beta -= H * (V if IS_BETA_HEADWISE else 1) + + if USE_INITIAL_STATE: + p_dh0 = dh0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_k[:, None] & mask_v[None, :]) + + tl.debug_barrier() + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + bos * H + i_h + p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_dq = dq + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + + if USE_INITIAL_STATE: + mask_h = mask_k[:, None] & mask_v[None, :] + p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_dk = tl.load(p_dk, mask=mask_k, other=0).to(tl.float32) + b_dv = tl.load(p_dv, mask=mask_v, other=0).to(tl.float32) + b_dk -= tl.sum(b_dv[None, :] * b_h, axis=1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + + b_h += b_k[:, None] * b_v[None, :] + b_dq = b_h * b_do[None, :] + d_q = tl.sum(b_dq, axis=1) * scale + tl.store(p_dq, d_q.to(p_dq.dtype.element_ty), mask=mask_k) + + p_k += H*K + p_v += H*V + p_do += H*V + p_dq += H*K + p_dk += H*K + p_dv += H*V + p_beta += H * (V if IS_BETA_HEADWISE else 1) + + +def fused_recurrent_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 1 + num_warps = 1 + + o = q.new_empty(NK, *v.shape) + if output_final_state: + final_state = q.new_empty(N, H, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NV, NK, N * H) + u = torch.empty_like(v) + fused_recurrent_delta_rule_fwd_kernel[grid]( + q, + k, + v, + u, + beta, + o, + initial_state, + final_state, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o, u, final_state + + +def fused_recurrent_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + dht: torch.Tensor, + do: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 1 + num_warps = 2 + + beta_vector = beta.ndim == v.ndim + + dq = q.new_empty(NV, *q.shape) + dk = q.new_empty(NV, *k.shape) + dv = q.new_empty(NK, *v.shape) + if beta_vector: + db = q.new_empty(NV, NK, B, T, H, V) + else: + db = q.new_empty(NV, B, T, H) + grid = (NV, NK, N * H) + + if initial_state is not None and initial_state.requires_grad: + dh0 = torch.empty_like(initial_state, dtype=torch.float32) + else: + dh0 = None + + fused_recurrent_delta_rule_bwd_kernel[grid]( + q, + k, + v, + beta, + initial_state, + dh0, + dht, + do, + dq, + dk, + dv, + db, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + NK=NK, + IS_BETA_HEADWISE=beta_vector, + num_warps=num_warps, + num_stages=num_stages, + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + db = db.sum((0, 1)) if beta_vector else db.sum(0) + + return dq, dk, dv, db, dh0 + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + else: + q_rstd, k_rstd = None, None + + o, u, final_state = fused_recurrent_delta_rule_fwd( + q=q, + k=k, + v=v, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + + ctx.save_for_backward(q, q_rstd, k, k_rstd, u, beta, initial_state) + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + q, q_rstd, k, k_rstd, v, beta, initial_state = ctx.saved_tensors + dq, dk, dv, db, dh0 = fused_recurrent_delta_rule_bwd( + q=q, + k=k, + v=v, + beta=beta, + dht=dht, + do=do, + scale=ctx.scale, + initial_state=initial_state, + cu_seqlens=ctx.cu_seqlens, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q), dk.to(k), dv.to(v), db.to(beta), None, dh0, None, None, None + + +@torch.compiler.disable +def fused_recurrent_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.delta_rule import fused_recurrent_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> beta = torch.rand(B, T, H, device='cuda').sigmoid() + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_recurrent_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = fused_recurrent_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + if beta is None: + beta = torch.ones_like(q[..., 0]) + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/delta_rule/naive.py b/fla/ops/delta_rule/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..a0cf44ef1756a42036bc69635c48cc581f5a3e68 --- /dev/null +++ b/fla/ops/delta_rule/naive.py @@ -0,0 +1,119 @@ + +import torch +from einops import rearrange + + +def delta_rule_recurrence(q, k, v, beta, initial_state=None, output_final_state=True): + orig_dtype = q.dtype + b, h, l, d_k = q.shape + q, k, v, beta = map(lambda x: x.float(), [q, k, v, beta]) + d_v = v.shape[-1] + o = torch.zeros_like(v) + S = torch.zeros(b, h, d_k, d_v).to(v) + q = q * (d_k ** -0.5) + + if beta.ndim < v.ndim: + beta = beta[..., None] + + if initial_state is not None: + S += initial_state + + for i in range(l): + _k = k[:, :, i] + _q = q[:, :, i] + _v = v[:, :, i].clone() + beta_i = beta[:, :, i] + _v = _v - (S.clone() * _k[..., None]).sum(-2) + _v = _v * beta_i + S = S.clone() + _k.unsqueeze(-1) * _v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) + S = None if output_final_state is False else S + return o.to(orig_dtype), S + + +def delta_rule_chunkwise(q, k, v, beta, chunk_size=32): + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * (d_k ** -0.5) + v = v * beta[..., None] + k_beta = k * beta[..., None] + + assert l % chunk_size == 0 + + # compute (I - tri(diag(beta) KK^T))^{-1} + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, k_beta]) + attn = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + + u = attn @ v + w = attn @ k_beta + S = k.new_zeros(b, h, d_k, d_v) + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i = q[:, :, i], k[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) + u_i = u[:, :, i] - w[:, :, i] @ S + o_inter = q_i @ S + o[:, :, i] = o_inter + attn @ u_i + S = S + k_i.transpose(-1, -2) @ u_i + + return rearrange(o, 'b h n c d -> b h (n c) d'), S + + +def delta_rule_parallel(q, k, v, beta, BM=128, BN=32): + b, h, l, d_k = q.shape + # d_v = v.shape[-1] + q = q * (d_k ** -0.5) + v = v * beta[..., None] + k_beta = k * beta[..., None] + # compute (I - tri(diag(beta) KK^T))^{-1} + q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BN), [q, k, v, k_beta]) + mask = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=0) + T = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, BN): + T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) + T = T + torch.eye(BN, dtype=torch.float, device=q.device) + + mask2 = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=1) + A_local = (q @ k.transpose(-1, -2)).masked_fill(mask2, 0) @ T + o_intra = A_local @ v + + # apply cumprod transition matrices on k to the last position within the chunk + k = k - ((k @ k.transpose(-1, -2)).masked_fill(mask, 0) @ T).transpose(-1, -2) @ k_beta + # apply cumprod transition matrices on q to the first position within the chunk + q = q - A_local @ k_beta + o_intra = A_local @ v + + A = torch.zeros(b, h, l, l, device=q.device) + + q, k, v, k_beta, o_intra = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, v, k_beta, o_intra]) + o = torch.empty_like(v) + for i in range(0, l, BM): + q_i = q[:, :, i:i+BM] + o_i = o_intra[:, :, i:i+BM] + # intra block + for j in range(i + BM - 2 * BN, i-BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + mask = torch.arange(i, i+BM) >= (j + BN) + A_ij = A_ij.masked_fill_(~mask[:, None].to(A_ij.device), 0) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + # inter block + for j in range(i - BN, -BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + o[:, :, i:i+BM] = o_i + + for i in range(0, l//BN): + A[:, :, i*BN:i*BN+BN, i*BN:i*BN+BN] = A_local[:, :, i] + + return o, A diff --git a/fla/ops/delta_rule/parallel.py b/fla/ops/delta_rule/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..903c1e631923c09fb85d203dffda02b3f59f602b --- /dev/null +++ b/fla/ops/delta_rule/parallel.py @@ -0,0 +1,403 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl +from einops import rearrange + +from fla.ops.delta_rule.wy_fast import fwd_prepare_T +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BT', 'K', 'V'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_transform_qk_fwd_kernel( + q, + k, + v, + beta, + o, + A, + q_new, + k_new, + A_local, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BT: tl.constexpr, + OUTPUT_ATTENTIONS: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(p_q.dtype.element_ty) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + p_T = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_T = tl.load(p_T, boundary_check=(0, 1)) + + o_i = tl.arange(0, BT) + m_t = o_i[:, None] >= o_i[None, :] + b_qk = tl.where(m_t, tl.dot(b_q, tl.trans(b_k), allow_tf32=False), 0).to(b_q.dtype) + m_t = o_i[:, None] > o_i[None, :] + b_kk = tl.where(m_t, tl.dot(b_k, tl.trans(b_k), allow_tf32=False), 0).to(b_k.dtype) + + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) + + b_qkT = tl.dot(b_qk, b_T, allow_tf32=False).to(b_k.dtype) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_a, b_qkT.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + b_kkT = tl.dot(b_kk, b_T, allow_tf32=False).to(b_k.dtype) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + tl.store(p_o, tl.dot(b_qkT, b_v).to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + p_q_new = tl.make_block_ptr(q_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_new, (b_q - tl.dot(b_qkT, b_k_beta, allow_tf32=False)).to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + + p_k_new = tl.make_block_ptr(k_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_k_new = b_k - tl.dot(tl.trans(b_kkT), b_k_beta, allow_tf32=False) + tl.store(p_k_new, b_k_new.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_transform_qk_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + chunk_size: int, + output_attentions: bool, +): + B, H, T, K = k.shape + BT = chunk_size + q_new = torch.empty_like(q) + k_new = torch.empty_like(k) + o = torch.empty_like(v) + grid = (triton.cdiv(T, BT), B*H) + V = v.shape[-1] + A_local = torch.empty_like(A) if output_attentions else None + chunk_transform_qk_fwd_kernel[grid]( + q, + k, + v, + beta, + o, + A, + q_new, + k_new, + A_local, + scale=scale, + T=T, + K=K, + V=V, + BT=BT, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + OUTPUT_ATTENTIONS=output_attentions, + ) + return q_new, k_new, o, A_local + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def save_intra_chunk_attn( + A, + A_local, + T, + BT: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + p_A = tl.make_block_ptr(A + i_bh * T * T, (T, T), (T, 1), (i_t * BT, i_t * BT), (BT, BT), (1, 0)) + p_A_local = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A_local = tl.load(p_A_local, boundary_check=(0, 1)) + tl.store(p_A, b_A_local.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_delta_rule_fwd_kernel( + q, + k, + k2, # original k + v, + beta, + o, + o_new, + attn, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + OUTPUT_ATTENTIONS: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [BT, BK] + b_q = tl.zeros([BT, BK], dtype=tl.float32) + b_q += tl.load(p_q, boundary_check=(0, 1)) + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + + # As opposed to Flashattention, this kernel requires scanning the KV blocks from right to left + # Q block and K block have overlap. + # masks required + for offset in range((i_t + 1) * BT - 2 * BS, i_t * BT - BS, -BS): + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) + p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS] + b_beta = tl.load(p_beta, boundary_check=(0,)) + # [BT, BS] + m_s = tl.arange(0, BT) >= (offset - i_t*BT + BS) + b_s = tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False) + b_s = tl.where(m_s[:, None], b_s, 0) + + b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + # Q block and K block have no overlap + # no need for mask, thereby saving flops + for offset in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) + p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) + + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS] + b_beta = tl.load(p_beta, boundary_check=(0,)) + # [BT, BS] + b_s = (tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False)) + # [BT, BV] + b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False).to(b_q.dtype) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + p_o_new = tl.make_block_ptr(o_new + i_bh * T*V, (T, V), (V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + tl.store(p_o_new, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +class ParallelDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, beta, scale, output_attentions): + B, H, T, K, V = *k.shape, v.shape[-1] + assert q.shape[-1] <= 128, 'The maximum supported sequence length is 128.' + BT, BS = 128, 32 + BK = triton.next_power_of_2(k.shape[-1]) + BV = triton.next_power_of_2(v.shape[-1]) + assert BT % BS == 0 + + A = fwd_prepare_T(k, beta, BS) + attn = q.new_zeros(B, H, T, T) if output_attentions else None + q_new, k_new, o, A_local = chunk_transform_qk_fwd( + q, + k, + v, + beta, + A, + scale, + BS, + output_attentions, + ) + + num_stages = 3 if K <= 64 else 2 + num_warps = 4 + grid = (triton.cdiv(T, BT), B * H) + o_new = torch.empty_like(o) + + parallel_delta_rule_fwd_kernel[grid]( + q=q_new, + k=k_new, + k2=k, + v=v, + beta=beta, + o=o, + o_new=o_new, + attn=attn, + T=T, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + num_stages=num_stages, + num_warps=num_warps, + ) + + if output_attentions: + grid = (triton.cdiv(T, BS), B * H) + save_intra_chunk_attn[grid]( + A=attn, + A_local=A_local, + T=T, + BT=BS, + ) + return o_new.to(q.dtype), attn + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, d_attn=None): + raise NotImplementedError('Backward pass is not implemented. Stay tuned!') + + +def parallel_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + output_attentions: bool = False, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + output_attentions (bool): + Whether to output the materialized attention scores of shape [B, H, T, T]. Default: `False`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + attn (torch.Tensor): + Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + o, attn = ParallelDeltaRuleFunction.apply(q, k, v, beta, scale, output_attentions) + return o, attn + + +def naive_delta_rule_parallel(q, k, v, beta, BM=128, BN=32): + b, h, l, d_k = q.shape + q = q * (d_k ** -0.5) + v = v * beta[..., None] + k_beta = k * beta[..., None] + # compute (I - tri(diag(beta) KK^T))^{-1} + q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BN), [q, k, v, k_beta]) + mask = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=0) + T = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, BN): + T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) + T = T + torch.eye(BN, dtype=q.dtype, device=q.device) + + mask2 = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=1) + A_local = (q @ k.transpose(-1, -2)).masked_fill(mask2, 0) @ T + o_intra = A_local @ v + + # apply cumprod transition matrices on k to the last position within the chunk + k = k - ((k @ k.transpose(-1, -2)).masked_fill(mask, 0) @ T).transpose(-1, -2) @ k_beta + # apply cumprod transition matrices on q to the first position within the chunk + q = q - A_local @ k_beta + o_intra = A_local @ v + + A = torch.zeros(b, h, l, l, device=q.device) + + q, k, v, k_beta, o_intra = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, v, k_beta, o_intra]) + o = torch.empty_like(v) + for i in range(0, l, BM): + q_i = q[:, :, i:i+BM] + o_i = o_intra[:, :, i:i+BM] + # intra block + for j in range(i + BM - 2 * BN, i-BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + mask = torch.arange(i, i+BM) >= (j + BN) + A_ij = A_ij.masked_fill_(~mask[:, None].to(A_ij.device), 0) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + # inter block + for j in range(i - BN, -BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + o[:, :, i:i+BM] = o_i + + for i in range(0, l//BN): + A[:, :, i*BN:i*BN+BN, i*BN:i*BN+BN] = A_local[:, :, i] + + return o, A diff --git a/fla/ops/delta_rule/wy_fast.py b/fla/ops/delta_rule/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..4119653053729fbc70925342413b23f878a38e70 --- /dev/null +++ b/fla/ops/delta_rule/wy_fast.py @@ -0,0 +1,301 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.solve_tril import solve_tril +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A.to(b_vb.dtype), b_vb, allow_tf32=False) + tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A.to(b_kb.dtype), b_kb, allow_tf32=False) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + A, + dw, + du, + dk, + dv, + dbeta, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_dbeta = tl.zeros([BT], dtype=tl.float32) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_v_beta), allow_tf32=False) + b_dv_beta = tl.dot(b_A, b_du, allow_tf32=False) + b_dv = b_dv_beta * b_beta[:, None] + b_dbeta += tl.sum(b_dv_beta * b_v, 1) + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dA += tl.dot(b_dw, tl.trans(b_k_beta), allow_tf32=False) + b_dk_beta = tl.dot(b_A, b_dw, allow_tf32=False) + b_dk = b_dk_beta * b_beta[:, None] + b_dbeta += tl.sum(b_dk_beta * b_k, 1) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dA, 0).to(k.dtype.element_ty) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)) + b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) + + b_dk_beta = tl.dot(b_dA, b_k, allow_tf32=False) + b_dbeta += tl.sum(b_dk_beta * b_k, 1) + b_dk += tl.dot(tl.trans(b_dA), b_k_beta, allow_tf32=False) + b_dk += b_dk_beta * b_beta[:, None] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + p_dbeta = tl.make_block_ptr(dbeta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + + +def prepare_wy_repr_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + A = chunk_scaled_dot_kkt_fwd( + k=k, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_size=64, + output_dtype=torch.float32, + chunk_indices=chunk_indices, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, A + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + u = torch.empty_like(v) + w = torch.empty_like(k) + recompute_w_u_fwd_kernel[(NT, B*H)]( + k, + v, + beta, + w, + u, + A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dbeta = torch.empty_like(beta) + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + k, + v, + beta, + A, + dw, + du, + dk, + dv, + dbeta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dk, dv, dbeta + + +fwd_prepare_wy_repr = prepare_wy_repr_fwd + +bwd_prepare_wy_repr = prepare_wy_repr_bwd + +fwd_recompute_w_u = recompute_w_u_fwd diff --git a/fla/ops/deltaformer/__init__.py b/fla/ops/deltaformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6d811ba8b19c41251677fc2a30b721f7f02a7dcc --- /dev/null +++ b/fla/ops/deltaformer/__init__.py @@ -0,0 +1,8 @@ + +from .naive import naive_deltaformer_attn +from .parallel import deltaformer_attn + +__all__ = [ + 'deltaformer_attn', + 'naive_deltaformer_attn', +] diff --git a/fla/ops/deltaformer/invcum.py b/fla/ops/deltaformer/invcum.py new file mode 100644 index 0000000000000000000000000000000000000000..2be2b295a06ca5e96fd05f2ff69a4dec65a64243 --- /dev/null +++ b/fla/ops/deltaformer/invcum.py @@ -0,0 +1,37 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch + + +def forward(u, w): + return torch.linalg.solve_triangular( + w.float(), + u.float(), + upper=False, + unitriangular=True, + ).to(u.dtype) + + +def forward_inplace(u, w): + u.copy_(forward(u, w)) + + +def backward_x(do, w): + return torch.linalg.solve_triangular( + w.tril(-1).mH.float(), + do.float(), + upper=True, + unitriangular=True, + ).to(do.dtype) + + +def backward(do, w, x): + du = torch.linalg.solve_triangular( + w.tril(-1).mH.float(), + do.float(), + upper=True, + unitriangular=True, + ).to(do.dtype) + dw = torch.bmm(-du, x.mH) + dw = dw.tril(-1) + return du, dw diff --git a/fla/ops/deltaformer/naive.py b/fla/ops/deltaformer/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..20bb89c6b4d18d828e8a413fbd908118b15e1dbd --- /dev/null +++ b/fla/ops/deltaformer/naive.py @@ -0,0 +1,150 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import math + +import torch + + +def tril_softmax(scores: torch.Tensor, strict: bool = True) -> torch.Tensor: + """ + Row-wise causal softmax over strictly lower-triangular (j < i) positions. + + Args: + scores: [B, H, T, T] raw attention scores (q @ k^T). + strict: if True, mask out diagonal as well (strictly causal). Otherwise include diagonal. + + Returns: + probs: [B, H, T, T] with probabilities on j < i (or j <= i if strict=False), zeros elsewhere. + """ + T = scores.size(-1) + device = scores.device + i = torch.arange(T, device=device).view(1, 1, T, 1) + j = torch.arange(T, device=device).view(1, 1, 1, T) + if strict: + mask = (j < i) + else: + mask = (j <= i) + + masked = scores.masked_fill(~mask, float('-inf')) + max_per_row = masked.max(dim=-1, keepdim=True).values + exp = (masked - max_per_row).exp() + exp = exp.masked_fill(~mask, 0.0) + denom = exp.sum(dim=-1, keepdim=True).clamp_min_(1e-20) + probs = exp / denom + return probs + + +def naive_causal_attention_bhtd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> torch.Tensor: + B, H, T, D = q.shape + qk_scale = 1.0 / math.sqrt(D) + scores = torch.matmul(q, k.transpose(-1, -2)) * qk_scale # [B, H, T, T] + causal_mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool() + scores = scores.masked_fill(causal_mask, float('-inf')) + attn_weights = torch.softmax(scores, dim=-1) # [B, H, T, T] + o = torch.matmul(attn_weights, v) # [B, H, T, D] + + return o + + +def naive_deltaformer_attn_head_first( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Naive reference implementation of DeltaFormer attention for head-first format. + + Two-stage process: + 1. Computes u[i] = v[i] - beta[i] * sum_{j torch.Tensor: + """ + Naive reference implementation of DeltaFormer attention for sequence-first format. + + Args: + q: [B, T, H, D] + k: [B, T, H, D] + v: [B, T, H, D] + beta: [B, T, H] or None (defaults to ones) + + Returns: + o: [B, T, H, D] + """ + assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "q,k,v must be [B,T,H,D]" + B, T, H, D = q.shape + assert k.shape == (B, T, H, D) and v.shape == (B, T, H, D) + + q_bhtd = q.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] + k_bhtd = k.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] + v_bhtd = v.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] + + if beta is not None: + assert beta.shape == (B, T, H) + beta_bhtd = beta.transpose(1, 2) # [B, T, H] -> [B, H, T] + else: + beta_bhtd = None + + o_bhtd = naive_deltaformer_attn_head_first(q_bhtd, k_bhtd, v_bhtd, beta_bhtd) + + o_bthd = o_bhtd.transpose(1, 2) # [B, H, T, D] -> [B, T, H, D] + + return o_bthd + + +__all__ = [ + 'naive_deltaformer_attn', + 'tril_softmax', +] diff --git a/fla/ops/deltaformer/parallel.py b/fla/ops/deltaformer/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..fd8f276edc1c5bbfd20a420bd1f43a178f2bb350 --- /dev/null +++ b/fla/ops/deltaformer/parallel.py @@ -0,0 +1,991 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import math +import warnings + +import torch +import triton +import triton.language as tl + +from . import invcum + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + +from fla.layers.utils import pad_input, unpad_input + +BLOCK_SIZE_C = 512 + + +def parallel_deltaformer_chunk_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + u: torch.Tensor, + qk_scale: float, + beta: torch.Tensor, +): + C, H, D = q.size() + T, _H, _D = k.size() + __C, __H = beta.size() + assert H == _H and D == _D and H == __H and __C == C + w = torch.empty(C, H, C, device=q.device, dtype=q.dtype) + lse = torch.empty(C, H, device=q.device, dtype=torch.float) + parallel_deltaformer_kernel(q, k, v, u, w, lse, qk_scale, beta) + return w, lse + + +def parallel_deltaformer_bwd_u_chunk( + q: torch.Tensor, + k: torch.Tensor, + lse: torch.Tensor, + grad_v: torch.Tensor, + fa_scale: float, + beta: torch.Tensor, +): + C, H, D = q.size() + T, _H, _D = k.size() + grad_u = torch.empty_like(q) + + def grid(META): + return (triton.cdiv(C, META['BLOCK_C']), H) + + parallel_deltaformer_bwd_kernel_u[grid]( + grad_u, q, k, grad_v, lse, beta, + H, T, C, D, fa_scale, + ) + return grad_u + + +def parallel_deltaformer_bwd_qk( + q: torch.Tensor, + k: torch.Tensor, + u: torch.Tensor, + lse: torch.Tensor, + grad_v: torch.Tensor, + qk_scale: float, + fa_scale: float, + beta: torch.Tensor, +): + T, H, D = k.size() + row_dot_sum = torch.empty_like(lse) + + def grid_bp(META): + return (triton.cdiv(T, META['BLOCK_C']), H) + + parallel_deltaformer_bwd_kernel_row_sum[grid_bp]( + row_dot_sum, q, k, grad_v, u, lse, + H, T, D, + fa_scale, + ) + grad_k = torch.empty_like(k) + grad_q = torch.empty_like(q) + + parallel_deltaformer_bwd_kernel_qk[grid_bp]( + grad_q, grad_k, q, k, grad_v, u, lse, beta, row_dot_sum, + H, T, D, + fa_scale, qk_scale, + ) + return grad_q, grad_k, row_dot_sum + + +def parallel_deltaformer_kernel( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + u: torch.Tensor, + w: torch.Tensor, + lse: torch.Tensor, + qk_scale: float, + beta: torch.Tensor, +) -> None: + C, H, D = q.size() + T, _H, _D = k.size() + + def grid(META): + return (triton.cdiv(C, META['BLOCK_C']), H) + + parallel_deltaformer_fwd_kernel[grid]( + q, k, v, u, w, lse, beta, + H, T, C, D, qk_scale, + ) + + +def _config_deltaformer(): + return [ + triton.Config({'BLOCK_C': BC, 'BLOCK_T': BT}, num_stages=ns, num_warps=nw) + for BC in [128, 64] + for BT in [64, 32] + for ns in [3, 2] + for nw in [8, 4] + ] + + +@triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) +@triton.jit +def parallel_deltaformer_fwd_kernel( + q_ptr, + k_ptr, + v_ptr, + u_ptr, + w_ptr, + lse_ptr, + beta_ptr, + H, + T, + C, + D: tl.constexpr, + qk_scale: float, + BLOCK_C: tl.constexpr, + BLOCK_T: tl.constexpr, +): + pid_c = tl.program_id(axis=0) + pid_h = tl.program_id(axis=1) + + rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C + colid_block = tl.arange(0, BLOCK_T) + + rowmax = tl.zeros([BLOCK_C], dtype=tl.float32) - float('inf') + rowsum = tl.zeros([BLOCK_C], dtype=tl.float32) + 1 + acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) + + q_blk_ptr = tl.make_block_ptr( + base=q_ptr + pid_h * D, + shape=(C, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + q = tl.load(q_blk_ptr, boundary_check=(0,)) + + for kv_i in range(0, T, BLOCK_T): + k_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_T), + order=(0, 1), + ) + k = tl.load(k_blk_ptr, boundary_check=(1,)) + qk = tl.dot(q, k) * qk_scale + + if kv_i >= T - C: + mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) + qk = tl.where(mask, -1e6, qk) + + rowmax_i = tl.maximum(rowmax, tl.max(qk, axis=1)) + qk -= rowmax_i[:, None] + p = tl.math.exp2(qk) + + rowsum_i = tl.sum(p, axis=1) + alpha = tl.math.exp2(rowmax - rowmax_i) + rowsum = rowsum * alpha + rowsum_i + acc = acc * alpha[:, None] + rowmax = rowmax_i + + if kv_i < T - C: + u_blk_ptr = tl.make_block_ptr( + base=u_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(kv_i, 0), + block_shape=(BLOCK_T, D), + order=(1, 0), + ) + u = tl.load(u_blk_ptr, boundary_check=(0,)) + acc = tl.dot(p.to(u_ptr.dtype.element_ty), u, acc) + + lse = rowmax + tl.math.log2(rowsum) + lse_block_ptr = lse_ptr + pid_h + rowid_block * H + lse_mask = rowid_block < C + tl.store(lse_block_ptr, lse, mask=lse_mask) + + v_ptr = tl.make_block_ptr( + base=v_ptr + pid_h * D, + shape=(C, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + acc = acc / rowsum[:, None] + + beta_ptr = tl.make_block_ptr( + base=beta_ptr + pid_h, + shape=(C,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + beta = tl.load(beta_ptr, boundary_check=(0,)) + acc = acc * beta[:, None] + + v = tl.load(v_ptr, boundary_check=(0,)) + u = v - acc.to(v_ptr.dtype.element_ty) + u_block_ptr = tl.make_block_ptr( + base=u_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(T - C + pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + tl.store(u_block_ptr, u, boundary_check=(0, 1)) + + for kv_i in range(T - C, T, BLOCK_T): + k_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_T), + order=(0, 1), + ) + k = tl.load(k_blk_ptr, boundary_check=(1,)) + qk = tl.dot(q, k) * qk_scale + + mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) + qk -= rowmax[:, None] + p = tl.math.exp2(qk) / rowsum[:, None] + p = tl.where(mask, 0, p) + w_blk_ptr = tl.make_block_ptr( + base=w_ptr + pid_h * C, + shape=(C, C), + strides=(H * C, 1), + offsets=(pid_c * BLOCK_C, kv_i - (T - C)), + block_shape=(BLOCK_C, BLOCK_T), + order=(1, 0), + ) + tl.store(w_blk_ptr, p.to(w_ptr.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) +@triton.jit +def parallel_deltaformer_bwd_kernel_u( + o_ptr, + q_ptr, + k_ptr, + v_ptr, + lse_ptr, + beta_ptr, + H, + T, + C, + D: tl.constexpr, + fa_scale, + BLOCK_C: tl.constexpr, + BLOCK_T: tl.constexpr, +): + pid_c = tl.program_id(axis=0) + pid_h = tl.program_id(axis=1) + + acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) + + q_blk_ptr = tl.make_block_ptr( + base=q_ptr + pid_h * D, + shape=(C, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + q = tl.load(q_blk_ptr, boundary_check=(0,)) + + for kv_i in range(0, T, BLOCK_T): + k_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_T), + order=(0, 1), + ) + k = tl.load(k_blk_ptr, boundary_check=(1,)) + qk = tl.dot(q, k) * fa_scale + + lse_blk_ptr = tl.make_block_ptr( + base=lse_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(kv_i,), + block_shape=(BLOCK_T,), + order=(0,), + ) + lse = tl.load(lse_blk_ptr, boundary_check=(0,)) + beta_blk_ptr = tl.make_block_ptr( + base=beta_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(kv_i,), + block_shape=(BLOCK_T,), + order=(0,), + ) + beta = tl.load(beta_blk_ptr, boundary_check=(0,)) + + p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] + + v_blk_ptr = tl.make_block_ptr( + base=v_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(kv_i, 0), + block_shape=(BLOCK_T, D), + order=(1, 0), + ) + v = tl.load(v_blk_ptr, boundary_check=(0,)) + acc = tl.dot(p.to(v_ptr.dtype.element_ty), v, acc) + + o_blk_ptr = tl.make_block_ptr( + base=o_ptr + pid_h * D, + shape=(C, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), boundary_check=(0,)) + + +@triton.autotune(configs=_config_deltaformer(), key=['T', 'D']) +@triton.jit +def parallel_deltaformer_bwd_kernel_row_sum( + row_dot_ptr, + q_ptr, + k_ptr, + grad_v_ptr, + u_ptr, + lse_ptr, + H, + T, + D: tl.constexpr, + fa_scale, + BLOCK_C: tl.constexpr, + BLOCK_T: tl.constexpr, +): + pid_c = tl.program_id(axis=0) + pid_h = tl.program_id(axis=1) + + rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C + colid_block = tl.arange(0, BLOCK_T) + + acc = tl.zeros([BLOCK_C], dtype=tl.float32) + + k_row_blk_ptr = tl.make_block_ptr( + base=q_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) + lse_blk_ptr = tl.make_block_ptr( + base=lse_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + lse = tl.load(lse_blk_ptr, boundary_check=(0,)) + grad_v_blk_ptr = tl.make_block_ptr( + base=grad_v_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) + + for kv_i in range(0, (pid_c + 1) * BLOCK_C, BLOCK_T): + k_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_T), + order=(0, 1), + ) + k = tl.load(k_blk_ptr, boundary_check=(1,)) + qk = tl.dot(k_row, k) * fa_scale + p = tl.math.exp2(qk - lse[:, None]) + + u_blk_ptr = tl.make_block_ptr( + base=u_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_T), + order=(0, 1), + ) + ut = tl.load(u_blk_ptr, boundary_check=(1,)) + dp = tl.dot(grad_v_row, ut) + if kv_i + BLOCK_T >= pid_c * BLOCK_C: + mask = (rowid_block[:, None] <= colid_block[None, :] + kv_i) + p = tl.where(mask, 0., p) + dp = tl.where(mask, 0., dp) + acc += tl.sum(p * dp, axis=1) + row_dot_block_ptr = tl.make_block_ptr( + base=row_dot_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + tl.store(row_dot_block_ptr, acc, boundary_check=(0,)) + + +@triton.autotune(configs=[triton.Config({'BLOCK_C': BC}, num_stages=ns, num_warps=nw) + for BC in [64, 32] + for ns in [4, 3] + for nw in [4]], key=['T', 'D']) +@triton.jit +def parallel_deltaformer_bwd_kernel_qk( + grad_q_ptr, + grad_k_ptr, + q_ptr, + k_ptr, + grad_v_ptr, + u_ptr, + lse_ptr, + beta_ptr, + row_dot_ptr, + H, + T, + D: tl.constexpr, + fa_scale: tl.constexpr, + qk_scale: tl.constexpr, + BLOCK_C: tl.constexpr, +): + pid_c = tl.program_id(axis=0) + pid_h = tl.program_id(axis=1) + block_i = tl.arange(0, BLOCK_C) + + acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) + + k_row_blk_ptr = tl.make_block_ptr( + base=q_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) + lse_blk_ptr = tl.make_block_ptr( + base=lse_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + lse = tl.load(lse_blk_ptr, boundary_check=(0,)) + beta_blk_ptr = tl.make_block_ptr( + base=beta_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + beta = tl.load(beta_blk_ptr, boundary_check=(0,)) + grad_v_blk_ptr = tl.make_block_ptr( + base=grad_v_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) + row_dot_blk_ptr = tl.make_block_ptr( + base=row_dot_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(pid_c * BLOCK_C,), + block_shape=(BLOCK_C,), + order=(0,), + ) + row_dot_row = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) + + for kv_i in range(0, pid_c * BLOCK_C, BLOCK_C): + k_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_C), + order=(0, 1), + ) + kt = tl.load(k_blk_ptr, boundary_check=(1,)) + qk = tl.dot(k_row, kt) * fa_scale + p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] + + u_blk_ptr = tl.make_block_ptr( + base=u_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_C), + order=(0, 1), + ) + ut = tl.load(u_blk_ptr) + dp = tl.dot(grad_v_row, ut) + da = p * (dp - row_dot_row[:, None]) + k = tl.trans(kt, 1, 0) + acc = tl.dot(da.to(k.dtype), k, acc) + + k_row_blk_ptr = tl.make_block_ptr( + base=k_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(pid_c * BLOCK_C, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + k_row_true = tl.load(k_row_blk_ptr, boundary_check=(0,)) + qk = tl.dot(k_row, tl.trans(k_row_true, 1, 0)) * fa_scale + p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] + u_blk_ptr = tl.make_block_ptr( + base=u_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, pid_c * BLOCK_C), + block_shape=(D, BLOCK_C), + order=(0, 1), + ) + ut = tl.load(u_blk_ptr) + dp = tl.dot(grad_v_row, ut) + dpm = dp - row_dot_row[:, None] + mask = block_i[None, :] < block_i[:, None] + p = tl.where(mask, p, 0.) + dpm = tl.where(mask, dpm, 0.) + da = p * dpm + daat = da + acc = tl.dot(daat.to(k_row.dtype), k_row_true, acc) + + grad_q_blk_ptr = tl.make_block_ptr( + base=grad_q_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(BLOCK_C * pid_c, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + acc = acc * qk_scale + tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), boundary_check=(0,)) + + daat = tl.trans(da, 1, 0) + acc = tl.dot(daat.to(k_row.dtype), k_row) + k_row = k_row_true + nu = -tl.trans(ut, 1, 0) + for kv_i in range((pid_c + 1) * BLOCK_C, T, BLOCK_C): + k_blk_ptr = tl.make_block_ptr( + base=q_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_C), + order=(0, 1), + ) + kt = tl.load(k_blk_ptr, boundary_check=(1,)) + lse_blk_ptr = tl.make_block_ptr( + base=lse_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(kv_i,), + block_shape=(BLOCK_C,), + order=(0,), + ) + lse = tl.load(lse_blk_ptr, boundary_check=(0,)) + beta_blk_ptr = tl.make_block_ptr( + base=beta_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(kv_i,), + block_shape=(BLOCK_C,), + order=(0,), + ) + beta = tl.load(beta_blk_ptr, boundary_check=(0,)) + qk = tl.dot(k_row, kt) * fa_scale + p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] + + grad_vt_blk_ptr = tl.make_block_ptr( + base=grad_v_ptr + pid_h * D, + shape=(D, T), + strides=(1, H * D), + offsets=(0, kv_i), + block_shape=(D, BLOCK_C), + order=(0, 1), + ) + grad_vt = tl.load(grad_vt_blk_ptr, boundary_check=(1,)) + row_dot_blk_ptr = tl.make_block_ptr( + base=row_dot_ptr + pid_h, + shape=(T,), + strides=(H,), + offsets=(kv_i,), + block_shape=(BLOCK_C,), + order=(0,), + ) + row_dot = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) + dp = tl.dot(nu, grad_vt) + da = p * (dp - row_dot[None, :]) + k = tl.trans(kt, 1, 0) + acc = tl.dot(da.to(k.dtype), k, acc) + + grad_k_blk_ptr = tl.make_block_ptr( + base=grad_k_ptr + pid_h * D, + shape=(T, D), + strides=(H * D, 1), + offsets=(BLOCK_C * pid_c, 0), + block_shape=(BLOCK_C, D), + order=(1, 0), + ) + acc = acc * qk_scale + tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), boundary_check=(0,)) + + +class ParallelDeltaformerFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qo: torch.Tensor, + ko: torch.Tensor, + vo: torch.Tensor, + betao: torch.Tensor | None = None, + C: int = BLOCK_SIZE_C, + cu_seqlens: torch.LongTensor | None = None, + ): + B, T, H, D = ko.size() + C = min(C, T) + ctx.C = C + ctx.cu_seqlens = cu_seqlens + + if cu_seqlens is not None: + need_aux = qo.requires_grad or ko.requires_grad or vo.requires_grad or (betao is not None and betao.requires_grad) + u, ws, lses = ParallelDeltaformerFunction._forward_impl( + qo, ko, vo, betao, C, need_aux=need_aux, cu_seqlens=cu_seqlens) + saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) + ctx.beta_is_none = betao is None + if need_aux: + ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) + else: + ctx.save_for_backward() + return u + + u, ws, lses = ParallelDeltaformerFunction._forward_impl(qo, ko, vo, betao, C, need_aux=True) + saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) + ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) + ctx.beta_is_none = betao is None + return u + + @staticmethod + def backward( + ctx, + grad_u: torch.Tensor, + ): + if getattr(ctx, 'cu_seqlens', None) is not None: + cu = ctx.cu_seqlens + qo, ko, vo, u_full, ws, lses, betao = ctx.saved_tensors + B, T_max, H, D = ko.size() + qk_scale = 1.0 / math.sqrt(D) + fa_scale = qk_scale / math.log(2) + + dq = torch.zeros_like(qo) + dk = torch.zeros_like(ko) + dv = torch.zeros_like(vo) + dbeta = None if ctx.beta_is_none else torch.zeros_like(betao) + + C = ctx.C + N = len(cu) - 1 + chunk_bases = [] + total = 0 + lengths = [] + for b in range(N): + L = int(cu[b + 1].item() - cu[b].item()) + lengths.append(L) + chunk_bases.append(total) + if L > 0: + total += (L + C - 1) // C + + for b in range(N): + L = lengths[b] + if L == 0: + continue + base = chunk_bases[b] + seq_start = int(cu[b].item()) + + seq_end = seq_start + L + q_seq = qo[0, seq_start:seq_end, :, :] + k_seq = ko[0, seq_start:seq_end, :, :] + u_seq = u_full[0, seq_start:seq_end, :, :] + beta_seq = betao[0, seq_start:seq_end, :] + lse_seq = lses[0, seq_start:seq_end, :] + go_seq = grad_u[0, seq_start:seq_end, :, :] + + gv_seq = torch.zeros_like(u_seq) + start = ((L - 1) // C) * C + for i_local in range(start, -1, -C): + Ci = min(C, L - i_local) + i0 = i_local + i1 = i_local + Ci + do = go_seq[i0:i1, :, :] + if i_local < L - C: + qi = k_seq[i0:i1, :, :] + ki = q_seq[i1:L, :, :] + lse_tail = lse_seq[i1:L, :] + beta_tail = beta_seq[i1:L, :] + du_tail = parallel_deltaformer_bwd_u_chunk(qi, ki, lse_tail, gv_seq[i1:L, :, :], fa_scale, beta_tail) + do = do - du_tail + Wpad = ws[base + (i_local // C)] + W = Wpad[:Ci, :, :Ci] + W_t = W.transpose(0, 1).contiguous() + du_chunk = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() + gv_seq[i0:i1, :, :].copy_(du_chunk) + + gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, gv_seq, qk_scale, fa_scale, beta_seq) + dq[0, seq_start:seq_end, :, :].copy_(gq) + dk[0, seq_start:seq_end, :, :].copy_(gk) + dv[0, seq_start:seq_end, :, :].copy_(gv_seq) + if dbeta is not None: + dbeta[0, seq_start:seq_end, :].copy_(gbeta) + + return dq, dk, dv, dbeta, None, None + qo, ko, vo, u, ws, lses, betao = ctx.saved_tensors + C = ctx.C + B, T, H, D = ko.size() + + grad_q = torch.zeros_like(qo) + grad_k = torch.zeros_like(ko) + grad_v = torch.zeros_like(vo) + grad_beta_out = None if ctx.beta_is_none else torch.zeros_like(betao) + + qk_scale = 1.0 / math.sqrt(D) + fa_scale = qk_scale / math.log(2) + + chunk_base = 0 + for b in range(B): + grad_v_seq = torch.empty(T, H, D, device=ko.device, dtype=ko.dtype) + for i in range(T - C, -1, -C): + Ci = min(C, T - i) + do = grad_u[b, i:i + Ci, :, :] + + if i < T - C: + qi = ko[b, i:i + Ci, :, :] + ki = qo[b, i + Ci:, :, :] + lse = lses[b, i + Ci:, :] + if not ctx.beta_is_none: + beta_single = betao[b, i + Ci:, :] + else: + beta_single = torch.ones(T - i - Ci, H, device=ko.device, dtype=ko.dtype) + du = parallel_deltaformer_bwd_u_chunk(qi, ki, lse, grad_v_seq[i + Ci:, :, :], fa_scale, beta_single) + do = grad_u[b, i:i + Ci, :, :] - du + + W = ws[chunk_base + (i // C)][:Ci, :, :Ci] + W_t = W.transpose(0, 1).contiguous() + du = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() + grad_v_seq[i:i + Ci, :, :].copy_(du) + + q_seq = qo[b] + k_seq = ko[b] + u_seq = u[b] + lse_seq = lses[b] + beta_seq = betao[b] if not ctx.beta_is_none else torch.ones(T, H, device=ko.device, dtype=ko.dtype) + + gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, grad_v_seq, qk_scale, fa_scale, beta_seq) + + grad_q[b].copy_(gq) + grad_k[b].copy_(gk) + grad_v[b].copy_(grad_v_seq) + if not ctx.beta_is_none: + grad_beta_out[b].copy_(gbeta) + + chunk_base += (T + C - 1) // C + + return grad_q, grad_k, grad_v, grad_beta_out, None, None + + @staticmethod + def _forward_impl( + qo: torch.Tensor, + ko: torch.Tensor, + vo: torch.Tensor, + betao: torch.Tensor | None, + C: int, + need_aux: bool, + cu_seqlens: torch.LongTensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T_max, H, D = ko.size() + C = min(C, T_max) + qk_scale = 1.0 / math.sqrt(D) + fa_scale = qk_scale / math.log(2) + + if cu_seqlens is None: + if betao is None: + beta_full = torch.ones(B, T_max, H, device=ko.device, dtype=ko.dtype) + else: + beta_full = betao + + u_full = torch.empty_like(vo) + if need_aux: + total_chunks = B * ((T_max + C - 1) // C) + ws = torch.empty(total_chunks, C, H, C, device=ko.device, dtype=ko.dtype) + lses = torch.empty(B, T_max, H, device=ko.device, dtype=torch.float) + chunk_base = 0 + else: + ws = None + lses = None + chunk_base = 0 + + for b in range(B): + for i in range(0, T_max, C): + Ci = min(C, T_max - i) + + qi = qo[b, i:i + Ci, :, :] + ki = ko[b, :i + Ci, :, :] + vi = vo[b, i:i + Ci, :, :] + ui_prev = u_full[b, :i + Ci, :, :] + betai = beta_full[b, i:i + Ci, :] + + w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) + w = w * betai.unsqueeze(-1) + if need_aux: + wpad = torch.zeros(C, H, C, device=ko.device, dtype=ko.dtype) + wpad[:Ci, :, :Ci].copy_(w) + ws[chunk_base + (i // C)].copy_(wpad) + lses[b, i:i + Ci, :].copy_(lse_chunk) + + u_chunk_view = u_full[b, i:i + Ci, :, :] + w_t = w.transpose(0, 1).contiguous() + u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() + invcum.forward_inplace(u_chunk_view_t, w_t) + u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) + + chunk_base += (T_max + C - 1) // C + + return u_full, ws, lses + + N = len(cu_seqlens) - 1 + assert cu_seqlens.dim() == 1 and cu_seqlens.size(0) == N + 1, "cu_seqlens must be [N+1]" + device = ko.device + dtype_k = ko.dtype + if betao is None: + beta_full = torch.ones(B, T_max, H, device=device, dtype=dtype_k) + else: + beta_full = betao + + u_full = torch.empty_like(vo) + if need_aux: + total_chunks = sum((max(0, int(cu_seqlens[b + 1].item() - cu_seqlens[b].item())) + C - 1) // C + for b in range(N)) + ws = torch.empty(total_chunks, C, H, C, device=device, dtype=dtype_k) + lses = torch.empty(B, T_max, H, device=device, dtype=torch.float) + chunk_base = 0 + else: + ws = None + lses = None + chunk_base = 0 + + for b in range(N): + seq_start = int(cu_seqlens[b].item()) + seq_end = int(cu_seqlens[b + 1].item()) + L = max(0, seq_end - seq_start) + if L == 0: + continue + + for i_local in range(0, L, C): + Ci = min(C, L - i_local) + li0 = i_local + li1 = i_local + Ci + + abs_start = seq_start + li0 + abs_end = seq_start + li1 + abs_context_end = seq_start + li1 + + qi = qo[0, abs_start:abs_end, :, :] + ki = ko[0, seq_start:abs_context_end, :, :] + vi = vo[0, abs_start:abs_end, :, :] + ui_prev = u_full[0, seq_start:abs_context_end, :, :] + betai = beta_full[0, abs_start:abs_end, :] + + w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) + w = w * betai.unsqueeze(-1) + if need_aux: + wpad = torch.zeros(C, H, C, device=device, dtype=dtype_k) + wpad[:Ci, :, :Ci].copy_(w) + ws[chunk_base + (i_local // C)].copy_(wpad) + lses[0, abs_start:abs_end, :].copy_(lse_chunk) + + u_chunk_view = u_full[0, abs_start:abs_end, :, :] + w_t = w.transpose(0, 1).contiguous() + u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() + invcum.forward_inplace(u_chunk_view_t, w_t) + u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) + + chunk_base += (L + C - 1) // C + + return u_full, ws, lses + + +def deltaformer_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + C: int = BLOCK_SIZE_C, +) -> torch.Tensor: + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + B, T, H, D = k.shape + C = min(C, T) + + u = ParallelDeltaformerFunction.apply(q, k, v, beta, C, cu_seqlens) + + if attention_mask is not None: + q_padded, (k_padded, u_padded), indices_q, cu_seqlens_lens, max_seq_lens = unpad_input(q, (k, u), attention_mask, T) + cu_seqlens_q, cu_seqlens_k = cu_seqlens_lens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q_padded, k_padded, u_padded, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1), + ) + o = pad_input(o, indices_q, B, T) + elif cu_seqlens is not None: + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), u.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1), + ).unsqueeze(0) + else: + o = flash_attn_func(q, k, u, causal=True, window_size=(-1, -1)) + + return o + + +__all__ = [ + 'deltaformer_attn', +] diff --git a/fla/ops/forgetting_attn/__init__.py b/fla/ops/forgetting_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..254c36ab18b0864a674a4f3ec54ea9213a8c810d --- /dev/null +++ b/fla/ops/forgetting_attn/__init__.py @@ -0,0 +1,8 @@ + +from .naive import naive_forgetting_attn +from .parallel import parallel_forgetting_attn + +__all__ = [ + 'naive_forgetting_attn', + 'parallel_forgetting_attn', +] diff --git a/fla/ops/forgetting_attn/naive.py b/fla/ops/forgetting_attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..665b2046416277679f8d195e21b2ed9ca7bafe50 --- /dev/null +++ b/fla/ops/forgetting_attn/naive.py @@ -0,0 +1,43 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + + +def naive_forgetting_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, +): + """ + Reference PyTorch implementation of forgetting attention. + + Args: + q: [B, T, HQ, D] + k: [B, T, H, D] + v: [B, T, H, D] + g: [B, T, HQ] + scale: float, optional. If None, defaults to 1 / sqrt(D) + + Returns: + output: [B, T, HQ, D] + """ + _, T, HQ, D = q.shape + H = k.shape[2] + G = HQ // H + + if scale is None: + scale = D ** -0.5 + + gc = g.float().cumsum(1) + mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=q.device)) + + ref = torch.einsum("bqhd,bkhd->bhqk", q.float() * scale, repeat(k, "b t h d -> b t (h g) d", g=G).float()) + ref = ref + rearrange(gc, "b t h -> b h t 1") - rearrange(gc, "b t h -> b h 1 t") + ref = ref.masked_fill(~mask.unsqueeze(0).unsqueeze(0), -float('inf')) + ref = torch.einsum("bhqk,bkhd->bqhd", F.softmax(ref, dim=-1), repeat(v, "b t h d -> b t (h g) d", g=G).float()) + + return ref diff --git a/fla/ops/forgetting_attn/parallel.py b/fla/ops/forgetting_attn/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..46ed9bae25a69c6c825581d38178e6348f23a7c7 --- /dev/null +++ b/fla/ops/forgetting_attn/parallel.py @@ -0,0 +1,62 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.attn.parallel import parallel_attn + + +def parallel_forgetting_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + GQA will be applied if HQ is divisible by H. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Log decay at rach time step (in **log space**) of shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + """ + + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + o = parallel_attn(q, k, v, g, scale, cu_seqlens) + return o diff --git a/fla/ops/gated_delta_product/__init__.py b/fla/ops/gated_delta_product/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ab6b529cb664b3092d820e08f7601ff0cc9df05 --- /dev/null +++ b/fla/ops/gated_delta_product/__init__.py @@ -0,0 +1,5 @@ +from .chunk import chunk_gated_delta_product + +__all__ = [ + "chunk_gated_delta_product", +] diff --git a/fla/ops/gated_delta_product/chunk.py b/fla/ops/gated_delta_product/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..6ca1a9cdf734d277c6328b5a07d3cee642cf4cd0 --- /dev/null +++ b/fla/ops/gated_delta_product/chunk.py @@ -0,0 +1,330 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +from einops import rearrange + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.delta_rule.chunk import chunk_delta_rule_bwd +from fla.ops.delta_rule.wy_fast import recompute_w_u_fwd as dn_recompute_w_u_fwd +from fla.ops.gated_delta_product.chunk_deltaproduct_h import chunk_gated_delta_product_fwd_h +from fla.ops.gated_delta_product.chunk_deltaproduct_o import chunk_gated_delta_product_fwd_o +from fla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule_bwd +from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd as gdn_recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum, solve_tril +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_gated_delta_product_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + num_householder: int = 1, + chunk_indices: torch.LongTensor | None = None, + chunk_indices_dp: torch.LongTensor | None = None, +): + cu_seqlens_dp = cu_seqlens * num_householder if cu_seqlens is not None else None + if g is not None: + g_interleaved = g.new_zeros(g.shape[0], g.shape[1], num_householder, g.shape[2], dtype=torch.float32) + g_interleaved[:, :, 0] = g + g_interleaved = rearrange(g_interleaved, 'b l n h -> b (l n) h').contiguous() + g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens, + output_dtype=torch.float32, chunk_indices=chunk_indices) + g_interleaved = chunk_local_cumsum( + g_interleaved, chunk_size=64, cu_seqlens=cu_seqlens_dp, output_dtype=torch.float32, chunk_indices=chunk_indices_dp + ) + else: + g_interleaved = None + g = None + # obtain WY representation. u is actually the new v. + A = chunk_scaled_dot_kkt_fwd( + k=k, + g=g_interleaved, + beta=beta, + cu_seqlens=cu_seqlens_dp, + output_dtype=torch.float32, + chunk_indices=chunk_indices_dp, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens_dp, + output_dtype=k.dtype, + chunk_indices=chunk_indices_dp, + ) + if g is not None: + w, u = gdn_recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g_interleaved, + cu_seqlens=cu_seqlens_dp, + chunk_indices=chunk_indices_dp, + ) + else: + w, u = dn_recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + cu_seqlens=cu_seqlens_dp, + chunk_indices=chunk_indices_dp, + ) + h, v_new, final_state = chunk_gated_delta_product_fwd_h( + k=k, + w=w, + u=u, + g=g_interleaved, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens_dp, + num_householder=num_householder, + chunk_indices=chunk_indices, + ) + o = chunk_gated_delta_product_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + num_householder=num_householder, + chunk_indices=chunk_indices, + ) + return g, g_interleaved, o, A, final_state + + +class ChunkGatedDeltaProductFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + num_householder: int, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + ): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + else: + q_rstd, k_rstd = None, None + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu + ) if cu_seqlens is not None else None + cu_seqlens_cpu_dp = cu_seqlens_cpu * num_householder if cu_seqlens_cpu is not None else None + chunk_indices_dp = prepare_chunk_indices( + cu_seqlens * num_householder, 64, cu_seqlens_cpu=cu_seqlens_cpu_dp + ) if cu_seqlens is not None else None + + g, g_interleaved, o, A, final_state = chunk_gated_delta_product_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + num_householder=num_householder, + chunk_indices=chunk_indices, + chunk_indices_dp=chunk_indices_dp, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g_interleaved, beta, A, initial_state, cu_seqlens, chunk_indices_dp) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.num_householder = num_householder + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens, chunk_indices_dp = ctx.saved_tensors + q_new = q.new_zeros(q.shape[0], q.shape[1], ctx.num_householder, q.shape[2], q.shape[3]) + q_new[:, :, -1] = q + do_new = do.new_zeros(do.shape[0], do.shape[1], ctx.num_householder, do.shape[2], do.shape[3]) + do_new[:, :, -1] = do + q_org, q = q, rearrange(q_new, 'b t n h d -> b (t n) h d') + do = rearrange(do_new, 'b t n h d -> b (t n) h d') + # call the gated deltanet kernel for now. + # TODO: optimize the backward pass like the forward pass. + if g is not None: + dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, + chunk_indices=chunk_indices_dp, + use_exp2=False, + ) + dg = rearrange(dg, 'b (l n) h -> b l n h ', n=ctx.num_householder)[:, :, 0].contiguous().to(g) + else: + dq, dk, dv, db, dh0 = chunk_delta_rule_bwd( + q=q, + k=k, + v=v, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, + chunk_indices=chunk_indices_dp, + ) + dg = None + dq = rearrange(dq, 'b (l n) h d -> b l n h d', n=ctx.num_householder)[:, :, -1].contiguous() + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q_org, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q), dk.to(k), dv.to(v), dg, db.to(beta), None, None, dh0, None, None, None, None + + +@torch.compiler.disable +def chunk_gated_delta_product( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + num_householder: int, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + num_householder (int): + Number of householder transformations to apply. Default: `1`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use qk l2norm within the kernel for saving GPU memory. + Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_product + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_product( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_product( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + assert q.dtype != torch.float32, "ChunkGatedDeltaProductFunction does not support float32. Please use bfloat16." + B, T, H, K, V = *q.shape, v.shape[-1] + assert k.shape == (B, T*num_householder, H, K) + assert v.shape == (B, T*num_householder, H, V) + assert beta.shape == (B, T*num_householder, H) + if g is not None: + assert g.shape == (B, T, H) + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaProductFunction.apply( + q, + k, + v, + g, + beta, + scale, + num_householder, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/gated_delta_product/chunk_deltaproduct_h.py b/fla/ops/gated_delta_product/chunk_deltaproduct_h.py new file mode 100644 index 0000000000000000000000000000000000000000..d7ba764335475ab5546d926edb274ccebf7eaca4 --- /dev/null +++ b/fla/ops/gated_delta_product/chunk_deltaproduct_h.py @@ -0,0 +1,505 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import IS_NVIDIA_HOPPER, USE_CUDA_GRAPH, autotune_cache_kwargs + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BV in [32, 64] + ], + key=['H', 'K', 'V', 'BT', 'USE_G'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_product_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + num_householder: tl.constexpr, # number of delta products + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * tl.cdiv(T // num_householder, BT) + + # [BK, BV] + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * H + i_h) * K*V + v += (bos * H + i_h) * V + k += (bos * H + i_h) * K + w += (bos * H + i_h) * K + if SAVE_NEW_VALUE: + v_new += (bos * H + i_h) * V + stride_v = H*V + stride_h = H*K*V + stride_k = H*K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K*V + if STORE_FINAL_STATE: + ht = ht + i_nh * K*V + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + if i_t % num_householder == 0: + i_t_true = i_t // num_householder + p_h1 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), + (BT, BV), (1, 0)) if SAVE_NEW_VALUE else None + b_v_new = tl.zeros([BT, BV], dtype=tl.float32) + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) + b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1)) + + if SAVE_NEW_VALUE: + p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v_new = b_v_new * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 = b_h1 * b_g_last + if K > 64: + b_h2 = b_h2 * b_g_last + if K > 128: + b_h3 = b_h3 * b_g_last + if K > 192: + b_h4 = b_h4 * b_g_last + b_v_new = b_v_new.to(k.dtype.element_ty) + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v_new) + if K > 64: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v_new) + if K > 128: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v_new) + if K > 192: + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v_new) + # epilogue + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [4, 3, 2] + for BV in [64, 32] + ], + key=['H', 'K', 'V', 'BT', 'BV', 'USE_G'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + dh += (boh * H + i_h) * K*V + dv += (bos * H + i_h) * V + dv2 += (bos * H + i_h) * V + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + w += (bos * H + i_h) * K + do += (bos * H + i_h) * V + stride_v = H*V + stride_h = H*K*V + stride_k = H*K + if USE_INITIAL_STATE: + dh0 += i_nh * K*V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K*V + + if USE_FINAL_STATE_GRADIENT: + p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + if K > 64: + p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + if K > 128: + p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + if K > 192: + p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + + for i_t in range(NT - 1, -1, -1): + p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + last_idx = min((i_t + 1) * BT, T) - 1 + bg_last = tl.load(g + (bos + last_idx) * H + i_h) + bg_last_exp = exp(bg_last) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_exp = exp(b_g) + else: + bg_last = None + last_idx = None + b_g = None + b_g_exp = None + + p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_wo = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_wo = tl.load(p_wo, boundary_check=(0, 1)) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + # Update dv + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dv += tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + + if K > 128: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + + if K > 192: + p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, boundary_check=(0, 1)) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + # Update dh + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + b_q = (b_q * scale).to(b_q.dtype) + b_dh1 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 64: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + b_q = (b_q * scale).to(b_q.dtype) + b_dh2 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 128: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + b_q = (b_q * scale).to(b_q.dtype) + b_dh3 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) + if K > 192: + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + b_q = (b_q * scale).to(b_q.dtype) + b_dh4 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_product_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, + num_householder: int = 1, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, u.shape[-1] + assert T % num_householder == 0, "T must be divisible by num_householder" + T_true = T // num_householder + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens // num_householder, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T_true, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - \ + 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens // num_householder, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + v_new = torch.empty_like(u) if save_new_value else None + + def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + chunk_gated_delta_product_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + num_householder=num_householder, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return h, v_new, final_state + + +def chunk_gated_delta_product_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + g: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor | None, + do: torch.Tensor, + dv: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *q.shape, do.shape[-1] + + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = 64 + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + dh = q.new_empty(B, NT, H, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + chunk_gated_delta_product_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dh, dh0, dv2 diff --git a/fla/ops/gated_delta_product/chunk_deltaproduct_o.py b/fla/ops/gated_delta_product/chunk_deltaproduct_o.py new file mode 100644 index 0000000000000000000000000000000000000000..0f57c2c4000a01ad76d215da3cb5299daa189080 --- /dev/null +++ b/fla/ops/gated_delta_product/chunk_deltaproduct_o.py @@ -0,0 +1,154 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + num_householder: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * num_householder * H + i_h) * K + v += (bos * num_householder * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K*V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_m = tl.where(m_A, exp(b_g[:, None] - b_g[None, :]), 0) + b_o = b_o * exp(b_g)[:, None] + else: + b_m = ((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)).to(tl.float32) + + for i_dp in range(num_householder): + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k+i_dp*H*K, (K, T), (1, num_householder*H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + b_A = b_A * b_m + p_v = tl.make_block_ptr(v+i_dp*H*V, (T, V), (H*V*num_householder, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_o += tl.dot(b_A.to(b_v.dtype), b_v) + b_o = b_o * scale + p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_product_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, # cumsum of log decay + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + num_householder: int = 1, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + assert q.shape[1] * num_householder == k.shape[1], "q.shape[1] * num_householder must be equal to k.shape[1]" + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + o = v.new_empty(B, T, H, V).fill_(-float('inf')) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_fwd_kernel_o[grid]( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T=T, + num_householder=num_householder, + H=H, + K=K, + V=V, + BT=BT, + ) + return o diff --git a/fla/ops/gated_delta_product/chunk_ref.py b/fla/ops/gated_delta_product/chunk_ref.py new file mode 100644 index 0000000000000000000000000000000000000000..d162035ef7df75f74394b395399bf0bc7958c3ee --- /dev/null +++ b/fla/ops/gated_delta_product/chunk_ref.py @@ -0,0 +1,66 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +from einops import rearrange + +from fla.ops.delta_rule import chunk_delta_rule +from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + +@torch.compiler.disable +def chunk_gated_delta_product_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + num_householder: int, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, +): + assert q.dtype != torch.float32, "ChunkGatedDeltaProductFunction does not support float32. Please use bfloat16." + B, T, H, K = q.shape + V = v.shape[-1] + assert k.shape == (B, T*num_householder, H, K) + assert v.shape == (B, T*num_householder, H, V) + assert beta.shape == (B, T*num_householder, H) + if g is not None: + assert g.shape == (B, T, H) + q_new = q.new_zeros(B, T, num_householder, H, K) + q_new[:, :, -1] = q + q = rearrange(q_new, 'b t n h d -> b (t n) h d') + + if g is not None: + g_new = g.new_zeros(B, T, num_householder, H, dtype=torch.float32) + g_new[:, :, 0] = g + g = rearrange(g_new, 'b t n h -> b (t n) h') + o, final_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens * num_householder if cu_seqlens is not None else None, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + scale=scale, + ) + else: + o, final_state = chunk_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens * num_householder if cu_seqlens is not None else None, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + scale=scale, + ) + o = rearrange(o, 'b (t n) h d -> b t n h d', n=num_householder) + return o[:, :, -1].contiguous(), final_state diff --git a/fla/ops/gated_delta_product/naive.py b/fla/ops/gated_delta_product/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..04b5b87f57b33c7842979a473f5ab4395b9e33d9 --- /dev/null +++ b/fla/ops/gated_delta_product/naive.py @@ -0,0 +1,36 @@ +import torch + + +def naive_recurrent_gated_delta_product(q, k, v, g, beta, scale, cu_seqlens, + initial_state=None, output_final_state=False, + num_householder=1): + q_original_dtype = q.dtype + B, T, H, K = q.shape + V = v.shape[-1] + assert k.shape == (B, T*num_householder, H, K) + assert v.shape == (B, T*num_householder, H, V) + assert beta.shape == (B, T*num_householder, H) + if g is not None: + assert g.shape == (B, T, H) + q, k, v, beta = map(lambda x: x.float(), (q, k, v, beta)) + + h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) + if initial_state is not None: + h = initial_state + + o = torch.zeros(B, T, H, V, dtype=torch.float32, device=q.device) + + for i in range(T): + if g is not None: + h = h * g[:, i, :].exp()[..., None, None] + # multiple state transition + for j in range(num_householder): + k_ij = k[:, i*num_householder+j, :, :] + v_ij = v[:, i*num_householder+j, :, :] + beta_ij = beta[:, i*num_householder+j, :] + h = h + (v_ij - (h * k_ij[..., None]).sum(-2)).unsqueeze(-2) * k_ij[..., None] * beta_ij[..., None, None] + # memory readout + q_i = q[:, i, :, :] + o_i = (h * q_i[..., None]).sum(-2) + o[:, i] = o_i + return o.to(q_original_dtype), h diff --git a/fla/ops/gated_delta_rule/__init__.py b/fla/ops/gated_delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f5f933971e2b71d21fff0cd7ce0ed954e4f29d85 --- /dev/null +++ b/fla/ops/gated_delta_rule/__init__.py @@ -0,0 +1,10 @@ +from .chunk import chunk_gated_delta_rule, chunk_gdn +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn +from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule + +__all__ = [ + "chunk_gated_delta_rule", "chunk_gdn", + "fused_recurrent_gated_delta_rule", "fused_recurrent_gdn", + "naive_chunk_gated_delta_rule", + "naive_recurrent_gated_delta_rule", +] diff --git a/fla/ops/gated_delta_rule/chunk.py b/fla/ops/gated_delta_rule/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..9980f5917f731b36aadf6f14bbbd4fa411f80e41 --- /dev/null +++ b/fla/ops/gated_delta_rule/chunk.py @@ -0,0 +1,436 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = True, + transpose_state_layout: bool = False, +): + g = chunk_local_cumsum( + g, + chunk_size=64, + scale=RCP_LN2 if use_exp2 else None, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + ) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + return g, o, A, final_state, initial_state + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = True, + transpose_state_layout: bool = False, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + ) + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + return dq, dk, dv, db, dg, dh0 + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + g, o, A, final_state, initial_state = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + transpose_state_layout=transpose_state_layout, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens, chunk_indices) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.cp_context = cp_context + ctx.transpose_state_layout = transpose_state_layout + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors + dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + transpose_state_layout=ctx.transpose_state_layout, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None, None, None + + +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + transpose_state_layout (Optional[bool]): + Whether to use the transposed state layout for the hidden state. + Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if 'head_first' in kwargs: + warnings.warn( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + use_qk_l2norm_in_kernel, + cp_context, + transpose_state_layout, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/fla/ops/gated_delta_rule/chunk_fwd.py b/fla/ops/gated_delta_rule/chunk_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..7c6c764f7788be7987248e8e4280a3df46bfc812 --- /dev/null +++ b/fla/ops/gated_delta_rule/chunk_fwd.py @@ -0,0 +1,402 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, exp2 +from fla.utils import IS_TF32_SUPPORTED, autotune_cache_kwargs + +if IS_TF32_SUPPORTED: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') +else: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=['H', 'K', 'BC'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_rule_fwd_kkt_solve_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass. + + This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel, + avoiding the HBM round-trip for the intermediate A matrix. + + Steps: + 1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers + 2. Apply gate and beta scaling + 3. Forward substitution on diagonal blocks + 4. Block merge to get full (I+A)^{-1} + 5. Write result to A (output) + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + k += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + + o_i = tl.arange(0, BC) + m_tc0 = (i_tc0 + o_i) < T + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + # load beta for each sub-chunk + p_b0 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,)) + p_b1 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,)) + p_b2 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,)) + p_b3 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,)) + b_b0 = tl.load(p_b0, boundary_check=(0,)).to(tl.float32) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + + # load gate if used + if USE_G: + p_g0 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,)) + p_g1 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,)) + p_g2 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,)) + p_g3 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,)) + + b_g0 = tl.load(p_g0, boundary_check=(0,)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0,)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0,)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0,)).to(tl.float32) + + ############################################################################ + # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T + ############################################################################ + + # 4 diagonal blocks + b_A00 = tl.zeros([BC, BC], dtype=tl.float32) + b_A11 = tl.zeros([BC, BC], dtype=tl.float32) + b_A22 = tl.zeros([BC, BC], dtype=tl.float32) + b_A33 = tl.zeros([BC, BC], dtype=tl.float32) + + # 6 off-diagonal blocks + b_A10 = tl.zeros([BC, BC], dtype=tl.float32) + b_A20 = tl.zeros([BC, BC], dtype=tl.float32) + b_A21 = tl.zeros([BC, BC], dtype=tl.float32) + b_A30 = tl.zeros([BC, BC], dtype=tl.float32) + b_A31 = tl.zeros([BC, BC], dtype=tl.float32) + b_A32 = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)) + # diagonal block 0 + b_A00 += tl.dot(b_k0, tl.trans(b_k0)) + + if i_tc1 < T: + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)) + # diagonal block 1 + b_A11 += tl.dot(b_k1, tl.trans(b_k1)) + # off-diagonal (1,0) + b_A10 += tl.dot(b_k1, tl.trans(b_k0)) + + if i_tc2 < T: + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)) + # diagonal block 2 + b_A22 += tl.dot(b_k2, tl.trans(b_k2)) + # off-diagonal (2,0), (2,1) + b_A20 += tl.dot(b_k2, tl.trans(b_k0)) + b_A21 += tl.dot(b_k2, tl.trans(b_k1)) + + if i_tc3 < T: + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)) + # diagonal block 3 + b_A33 += tl.dot(b_k3, tl.trans(b_k3)) + # off-diagonal (3,0), (3,1), (3,2) + b_A30 += tl.dot(b_k3, tl.trans(b_k0)) + b_A31 += tl.dot(b_k3, tl.trans(b_k1)) + b_A32 += tl.dot(b_k3, tl.trans(b_k2)) + + ############################################################################ + # Step 2: apply gate and beta scaling + ############################################################################ + + # apply gate, beta scaling, and masking + # m_d: strictly lower triangular mask for diagonal blocks + # m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when + # out-of-bounds g loads as 0 via boundary_check and exp(0 - g_inbounds) overflows + m_d = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + if USE_G: + if USE_EXP2: + b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.) + b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.) + b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.) + b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.) + + b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.) + b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.) + b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.) + b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.) + b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.) + b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.) + else: + b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp(b_g0[:, None] - b_g0[None, :]), 0.) + b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp(b_g1[:, None] - b_g1[None, :]), 0.) + b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp(b_g2[:, None] - b_g2[None, :]), 0.) + b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp(b_g3[:, None] - b_g3[None, :]), 0.) + + b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp(b_g1[:, None] - b_g0[None, :]), 0.) + b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp(b_g2[:, None] - b_g0[None, :]), 0.) + b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp(b_g2[:, None] - b_g1[None, :]), 0.) + b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp(b_g3[:, None] - b_g0[None, :]), 0.) + b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp(b_g3[:, None] - b_g1[None, :]), 0.) + b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp(b_g3[:, None] - b_g2[None, :]), 0.) + else: + b_A00 = tl.where(m_d, b_A00, 0.) + b_A11 = tl.where(m_d, b_A11, 0.) + b_A22 = tl.where(m_d, b_A22, 0.) + b_A33 = tl.where(m_d, b_A33, 0.) + + # diagonal blocks: scaled by beta + b_A00 = b_A00 * b_b0[:, None] + b_A11 = b_A11 * b_b1[:, None] + b_A22 = b_A22 * b_b2[:, None] + b_A33 = b_A33 * b_b3[:, None] + + # off-diagonal blocks: full block, scaled by beta + b_A10 = b_A10 * b_b1[:, None] + b_A20 = b_A20 * b_b2[:, None] + b_A21 = b_A21 * b_b2[:, None] + b_A30 = b_A30 * b_b3[:, None] + b_A31 = b_A31 * b_b3[:, None] + b_A32 = b_A32 * b_b3[:, None] + + ############################################################################ + # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1} + # + # Same algorithm as solve_tril, but rows are extracted from in-register + # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of + # tl.load from HBM. + ############################################################################ + + b_Ai00 = -b_A00 + b_Ai11 = -b_A11 + b_Ai22 = -b_A22 + b_Ai33 = -b_A33 + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(2, min(BC, T - i_tc1)): + b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0) + b_a11 = tl.where(o_i < i, b_a11, 0.) + b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11) + for i in range(2, min(BC, T - i_tc2)): + b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0) + b_a22 = tl.where(o_i < i, b_a22, 0.) + b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22) + for i in range(2, min(BC, T - i_tc3)): + b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0) + b_a33 = tl.where(o_i < i, b_a33, 0.) + b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ############################################################################ + # Step 4: block merge -> full (I + A)^{-1} + ############################################################################ + + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ############################################################################ + # Step 5: store full (I + A)^{-1} to output A + ############################################################################ + + p_A00 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_A10 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_A11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_A20 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_A21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_A22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_A30 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_A31 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_A32 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_A33 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_intra( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r""" + GDN intra-chunk forward: fused kkt + solve_tril + recompute_w_u. + + Equivalent to: + A = chunk_scaled_dot_kkt_fwd(k, g, beta, ...) # kernel 1 + A = solve_tril(A, ...) # kernel 2 + w, u = recompute_w_u_fwd(k, v, beta, A, g, ...) # kernel 3 + + Fuses kernels 1+2 into a single kernel, reducing from 3 to 2 kernel launches + and eliminating the HBM round-trip for the intermediate A matrix. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + v (torch.Tensor): + The value tensor of shape `[B, T, H, V]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths. Default: `None`. + chunk_size (int): + The chunk size. Default: 64. + chunk_indices (torch.LongTensor): + Precomputed chunk indices. Default: `None`. + + Returns: + w (torch.Tensor): shape `[B, T, H, K]` + u (torch.Tensor): shape `[B, T, H, V]` + A (torch.Tensor): shape `[B, T, H, BT]`, the solved (I+A)^{-1} matrix + """ + B, T, H, K = k.shape + BT = chunk_size + BC = 16 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + # Step 1: fused kkt + solve_tril + A = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * H)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + USE_EXP2=use_exp2, + ) + + # Step 2: recompute_w_u + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=use_exp2, + ) + return w, u, A diff --git a/fla/ops/gated_delta_rule/fused_recurrent.py b/fla/ops/gated_delta_rule/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..3f470a3ab035002ddce74580c16c6e6c78b21383 --- /dev/null +++ b/fla/ops/gated_delta_rule/fused_recurrent.py @@ -0,0 +1,406 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, exp2 +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + gk, + gv, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + USE_EXP2: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if USE_G: + p_g = g + bos * HV + i_hv + if USE_GK: + p_gk = gk + (bos * HV + i_hv) * K + o_k + if USE_GV: + p_gv = gv + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + bos * HV + i_hv + else: + p_beta = beta + (bos * HV + i_hv) * V + o_v + + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + if TRANSPOSE_STATE: + mask_h = mask_v[:, None] & mask_k[None, :] + else: + mask_h = mask_k[:, None] & mask_v[None, :] + + if TRANSPOSE_STATE: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if TRANSPOSE_STATE: + p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in tl.range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta).to(tl.float32) + else: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + if USE_EXP2: + b_h *= exp2(b_g) + else: + b_h *= exp(b_g) + + if USE_GK: + b_gk = tl.load(p_gk).to(tl.float32) + if USE_EXP2: + if TRANSPOSE_STATE: + b_h *= exp2(b_gk[None, :]) + else: + b_h *= exp2(b_gk[:, None]) + else: + if TRANSPOSE_STATE: + b_h *= exp(b_gk[None, :]) + else: + b_h *= exp(b_gk[:, None]) + + if USE_GV: + b_gv = tl.load(p_gv).to(tl.float32) + if USE_EXP2: + if TRANSPOSE_STATE: + b_h *= exp2(b_gv[:, None]) + else: + b_h *= exp2(b_gv[None, :]) + else: + if TRANSPOSE_STATE: + b_h *= exp(b_gv[:, None]) + else: + b_h *= exp(b_gv[None, :]) + + if TRANSPOSE_STATE: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1)) + b_h += b_v[:, None] * b_k[None, :] + b_o = tl.sum(b_h * b_q[None, :], 1) + else: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) + b_h += b_k[:, None] * b_v + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_v += HV*V + if USE_G: + p_g += HV + if USE_GK: + p_gk += HV*K + if USE_GV: + p_gv += HV*V + p_beta += HV * (1 if IS_BETA_HEADWISE else V) + p_o += HV*V + + if STORE_FINAL_STATE: + if TRANSPOSE_STATE: + p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V) + NV = triton.cdiv(V, BV) + + o = torch.empty_like(v) + if output_final_state: + if transpose_state_layout: + final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) + else: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NV, N * HV) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim != v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + num_warps=1, + num_stages=3, + ) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, + ): + o, final_state = fused_recurrent_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + use_exp2=use_exp2, + transpose_state_layout=transpose_state_layout, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps.", + ) + + +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. Default: `None`. + gk (torch.Tensor): + gk (decays) of shape `[B, T, HV, K]`. Default: `None`. + gv (torch.Tensor): + gv (decays) of shape `[B, T, HV, V]`. Default: `None`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + transpose_state_layout (bool): + Whether to use transposed state layout `[V, K]` instead of `[K, V]`. Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if beta is None: + beta = torch.ones_like(q[..., 0]) + + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + g, + gk, + gv, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + use_exp2, + transpose_state_layout, + ) + return o, final_state + + +fused_recurrent_gdn = fused_recurrent_gated_delta_rule diff --git a/fla/ops/gated_delta_rule/naive.py b/fla/ops/gated_delta_rule/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..0747e9bca0d10ae98f684e4af2d430f16238d2fc --- /dev/null +++ b/fla/ops/gated_delta_rule/naive.py @@ -0,0 +1,156 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of recurrent gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + beta: [B, T, H] + g: [B, T, H] + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def naive_chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of chunk gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + g: [B, T, H] + beta: [B, T, H] + chunk_size: int + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + + q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g]) + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + k_beta = k * beta[..., None] + assert l % chunk_size == 0 + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, k_beta, decay = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, k_beta, decay.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) + decay_exp = decay.exp()[..., None] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + attn = attn + k_cumsum = attn @ v + k_cumdecay = attn @ (k_beta * decay_exp) + v = k_cumsum + + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S = initial_state.to(torch.float32) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..52b2c9249144efe80a2922382fe4f82d2a0d11c5 --- /dev/null +++ b/fla/ops/gated_delta_rule/wy_fast.py @@ -0,0 +1,348 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, exp2 +from fla.utils import IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem + +if IS_NVIDIA_BLACKWELL: + """ + Compute tl.dot with SM100 workaround. + + On SM100 (Blackwell) GPUs, wraps the result in inline assembly to prevent + the TritonGPUHoistTMEMAlloc pass from incorrectly fusing add and dot operations. + See: https://github.com/fla-org/flash-linear-attention/issues/638 + + TODO: Remove this workaround once the Triton compiler bug is fixed. + Track upstream issue at: https://github.com/triton-lang/triton/issues/8695 + """ + @triton.jit + def safe_dot(a, b): + return tl.inline_asm_elementwise( + asm="mov.f32 $0, $1;", + constraints="=r,r", + args=[tl.dot(a, b)], + dtype=tl.float32, + is_pure=True, + pack=1, + ) +else: + @triton.jit + def safe_dot(a, b): + return tl.dot(a, b) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + if USE_EXP2: + b_g = exp2(tl.load(p_g, boundary_check=(0,))) + else: + b_g = exp(tl.load(p_g, boundary_check=(0,))) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + if USE_G: + b_kb *= b_g[:, None] + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + g, + A, + dw, + du, + dk, + dv, + db, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_b = tl.load(p_b, boundary_check=(0,)) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + if USE_G: + p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + if USE_EXP2: + b_g_exp = exp2(b_g) + else: + b_g_exp = tl.exp(b_g) + b_dg = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + if USE_G: + b_kbg = b_k * (b_b * b_g_exp)[:, None] + else: + b_kbg = b_k * b_b[:, None] + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + if USE_G: + b_dk = b_dkbg * (b_g_exp * b_b)[:, None] + b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) + b_dg += tl.sum(b_dkbg * b_kbg, 1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + if USE_G: + if USE_EXP2: + b_dA *= exp2(b_g[:, None] - b_g[None, :]) + else: + b_dA *= exp(b_g[:, None] - b_g[None, :]) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) + + tl.debug_barrier() + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kt = tl.trans(b_k) + b_kb = b_k * b_b[:, None] + + b_A += tl.dot(b_k, b_kt) + b_dkb = tl.dot(b_dA, b_k) + b_db += tl.sum(b_dkb * b_k, 1) + b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA)) + b_dk += tl.load(p_dk, boundary_check=(0, 1)) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + b_A *= b_b[:, None] + if USE_G: + b_AdA = b_dA * b_A + p_dg = tl.make_block_ptr(dg + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + recompute_w_u_fwd_kernel[(NT, B*H)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + ) + return w, u + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + g: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dg = torch.empty_like(g) if g is not None else None + db = torch.empty_like(beta) + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=du, + dk=dk, + dv=dv, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + USE_EXP2=use_exp2, + ) + return dk, dv, db, dg + + +fwd_recompute_w_u = recompute_w_u_fwd +bwd_prepare_wy_repr = prepare_wy_repr_bwd diff --git a/fla/ops/gated_oja_rule/__init__.py b/fla/ops/gated_oja_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4bb7a799dfc69279d7e29d37896d78e9b296b421 --- /dev/null +++ b/fla/ops/gated_oja_rule/__init__.py @@ -0,0 +1,7 @@ +from .chunk import chunk_gated_oja_rule +from .fused_recurrent import fused_recurrent_gated_oja_rule + +__all__ = [ + "chunk_gated_oja_rule", + "fused_recurrent_gated_oja_rule" +] diff --git a/fla/ops/gated_oja_rule/chunk.py b/fla/ops/gated_oja_rule/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..4f578b3d4430c301202800266e7983452858e50d --- /dev/null +++ b/fla/ops/gated_oja_rule/chunk.py @@ -0,0 +1,338 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.gated_oja_rule.chunk_h import chunk_oja_bwd_dhu, chunk_oja_bwd_dvwg_h, chunk_oja_fwd_h +from fla.ops.gated_oja_rule.chunk_kkt import chunk_scaled_dot_kkt_bwd_gk, chunk_scaled_dot_kkt_fwd +from fla.ops.gated_oja_rule.chunk_o import ( + chunk_oja_bwd_dA, + chunk_oja_bwd_dqk, + chunk_oja_bwd_dv_o, + chunk_oja_fwd_o, +) +from fla.ops.gated_oja_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum, solve_tril +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_oja_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + g_cumsum: bool = True, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + if g_cumsum: + gv = chunk_local_cumsum(gv, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + A = chunk_scaled_dot_kkt_fwd( + k=v, + gk=gv, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=torch.float32 + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype + ) + w, u, vg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gv=gv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + h, k_new, final_state = chunk_oja_fwd_h( + v=vg, + w=w, + u=u, + gv=gv, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + _, o = chunk_oja_fwd_o( + q=q, + k=k_new, + v=v, + h=h, + gv=gv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return gv, o, A, final_state + + +def chunk_oja_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + o: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dgk: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + w, u, vg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gv=gv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + h, k_new, _ = chunk_oja_fwd_h( + v=vg, + w=w, + u=u, + gv=gv, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dAqk = chunk_oja_bwd_dA( + v=v, + gv=gv, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + Aqk, dq, dk_new = chunk_oja_bwd_dqk( + q=q, + k=k_new, + h=h, + gv=gv, + dA=dAqk, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dh, dh0, dk_new = chunk_oja_bwd_dhu( + q=q, + vg=vg, + w=w, + gv=gv, + h0=initial_state, + dht=dht, + do=do, + dk=dk_new, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + states_in_fp32=False, + ) + + dv, dw, dgv_last = chunk_oja_bwd_dvwg_h( + k=k_new, + v=v, + gv=gv, + h=h, + dh=dh, + dk=dk_new, + dgk=dgk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dv, dgv1 = chunk_oja_bwd_dv_o( + v=v, + gv=gv, + o=o, + A=Aqk, + dv=dv, + do=do, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dk, dv1, db, dgv2, dAvv = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + gv=gv, + A=A, + dw=dw, + du=dk_new, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dv2, dgv3, db2 = chunk_scaled_dot_kkt_bwd_gk( + k=v, + g=gv, + beta=beta, + dA=dAvv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dv = dv.add_(dv1).add_(dv2) + db = db.add_(db2) + dgv = dgv_last.add_(chunk_local_cumsum( + dgv1.add_(dgv2).add_(dgv3), chunk_size=64, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + )) + return dq, dk, dv, db, dgv, dh0 + + +class ChunkOJAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + use_q_l2norm: bool = False, + use_k_l2norm: bool = False, + ): + q_rstd, k_rstd = None, None + if use_q_l2norm: + q, q_rstd = l2norm_fwd(q) + if use_k_l2norm: + k, k_rstd = l2norm_fwd(k) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + gv, o, A, final_state = chunk_oja_fwd( + q=q, + k=k, + v=v, + gv=gv, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, gv, beta, A, o, initial_state, cu_seqlens, chunk_indices) + ctx.scale = scale + ctx.use_q_l2norm = use_q_l2norm + ctx.use_k_l2norm = use_k_l2norm + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor + ): + q, q_rstd, k, k_rstd, v, gv, beta, A, o, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors + dq, dk, dv, db, dg, dh0 = chunk_oja_bwd( + q=q, + k=k, + v=v, + gv=gv, + beta=beta, + A=A, + o=o, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + if ctx.use_q_l2norm: + dq = l2norm_bwd(q, q_rstd, dq) + if ctx.use_k_l2norm: + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q), dk.to(k), dv.to(v), dg.to(gv), db.to(beta), None, dh0, None, None, None, None, None + + +@torch.compiler.disable +def chunk_gated_oja_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_q_l2norm: bool = False, + use_k_l2norm: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + **kwargs, +): + if 'head_first' in kwargs: + warnings.warn( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead." + ) + if 'use_qk_l2norm_in_kernel' in kwargs and (not use_q_l2norm and not use_k_l2norm): + use_q_l2norm = True + use_k_l2norm = True + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkOJAFunction.apply( + q, + k, + v, + gv, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + use_q_l2norm, + use_k_l2norm + ) + return o, final_state diff --git a/fla/ops/gated_oja_rule/chunk_h.py b/fla/ops/gated_oja_rule/chunk_h.py new file mode 100644 index 0000000000000000000000000000000000000000..c0e9839fa851a52fac32276378c0df15a2e7421a --- /dev/null +++ b/fla/ops/gated_oja_rule/chunk_h.py @@ -0,0 +1,807 @@ + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import check_shared_mem, is_nvidia_hopper, use_cuda_graph + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] + + +@triton.heuristics({ + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'SAVE_NEW_KEY': lambda args: args['k_new'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BK in [32, 64] + ], + key=['H', 'K', 'V', 'BT'], + use_cuda_graph=use_cuda_graph, +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_fwd_kernel_h_blockdim64( + v, + u, + w, + k_new, + gv, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_KEY: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # (triton.cdiv(K, meta['BK']), N*H) + i_k, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h1 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 64: + b_h2 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 128: + b_h3 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 192: + b_h4 = tl.zeros([BK, 64], dtype=tl.float32) + + # calculate offset + h += ((boh * H + i_h) * K*V).to(tl.int64) + v += ((bos * H + i_h) * V).to(tl.int64) + u += ((bos * H + i_h) * K).to(tl.int64) + w += ((bos * H + i_h) * V).to(tl.int64) + if SAVE_NEW_KEY: + k_new += ((bos * H + i_h) * K).to(tl.int64) + stride_v = H*V + stride_h = H*K*V + stride_k = H*K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K*V + if STORE_FINAL_STATE: + ht = ht + i_nh * K*V + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if V > 64: + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if V > 128: + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if V > 192: + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + p_h1 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if V > 64: + p_h2 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if V > 128: + p_h3 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if V > 192: + p_h4 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_k = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) # BT BK + if V > 64: + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_k += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + if V > 128: + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_k += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + if V > 192: + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_k += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + + p_u = tl.make_block_ptr(u, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_u, boundary_check=(0, 1)) - b_k + + if SAVE_NEW_KEY: + p_k = tl.make_block_ptr(k_new, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_k, b_k.to(p_k.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + + if USE_GV: + o_v1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v1, mask=(o_v1 < V), other=0.) + b_h1 *= exp(b_gk_last1)[None, :] + if V > 64: + o_v2 = 64 + o_v1 + b_gk_last2 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v2, mask=(o_v2 < V), other=0.) + b_h2 *= exp(b_gk_last2)[None, :] + if V > 128: + o_v3 = 128 + o_v1 + b_gk_last3 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v3, mask=(o_v3 < V), other=0.) + b_h3 *= exp(b_gk_last3)[None, :] + if V > 192: + o_v4 = 192 + o_v1 + b_gk_last4 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v4, mask=(o_v4 < V), other=0.) + b_h4 *= exp(b_gk_last4)[None, :] + + b_k = b_k.to(v.dtype.element_ty) # BT BK + + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) # BT BV + b_h1 += tl.dot(tl.trans(b_k), b_v) + if V > 64: + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_h2 += tl.dot(tl.trans(b_k), b_v) + if V > 128: + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_h3 += tl.dot(tl.trans(b_k), b_v) + if V > 192: + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_h4 += tl.dot(tl.trans(b_k), b_v) + # epilogue + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if V > 64: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if V > 128: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if V > 192: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_fwd_h( + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + gv: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + save_new_key: bool = True, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T, H, V, K = *v.shape, u.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + assert V <= 256, "current kernel does not support head dimension larger than 256." + + h = v.new_empty(B, NT, H, K, V) + final_state = v.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + + k_new = torch.empty_like(u) if save_new_key else None + def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) + chunk_oja_fwd_kernel_h_blockdim64[grid]( + v=v, + u=u, + w=w, + k_new=k_new, + gv=gv, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT + ) + return h, k_new, final_state + + +@triton.heuristics({ + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [4, 3, 2] + for BK in [64, 32] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'USE_GV'], + use_cuda_graph=use_cuda_graph, +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_bwd_kernel_dhu_blockdim64( + q, + vg, + w, + gv, + dht, + dh0, + do, + dh, + dk, + dk2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + USE_GV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_k, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh1 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 64: + b_dh2 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 128: + b_dh3 = tl.zeros([BK, 64], dtype=tl.float32) + if V > 192: + b_dh4 = tl.zeros([BK, 64], dtype=tl.float32) + + # calculate offset + q += ((bos * H + i_h) * K).to(tl.int64) + vg += ((bos * H + i_h) * V).to(tl.int64) + w += ((bos * H + i_h) * V).to(tl.int64) + do += ((bos * H + i_h) * V).to(tl.int64) + dk += ((bos * H + i_h) * K).to(tl.int64) + dk2 += ((bos * H + i_h) * K).to(tl.int64) + dh += ((boh * H + i_h) * K*V).to(tl.int64) + if USE_GV: + gv += ((bos * H + i_h) * V).to(tl.int64) + + stride_v = H*V + stride_h = H*K*V + stride_k = H*K + if USE_INITIAL_STATE: + dh0 += i_nh * K*V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K*V + + if USE_FINAL_STATE_GRADIENT: + p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) # [BK, BV] + b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + if V > 64: + p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + if V > 128: + p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + if V > 192: + p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + + for i_t in range(NT - 1, -1, -1): + p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if V > 64: + p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if V > 128: + p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + if V > 192: + p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + + # Update dk_new, 按K切分 + p_dk = tl.make_block_ptr(dk, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] + p_dk2 = tl.make_block_ptr(dk2, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] + + if V > 0: + p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) # [BT, BV] + b_dk = tl.dot(b_v, tl.trans(b_dh1).to(b_v.dtype)) # [BT, BV] @ [BV, BK] -> [BT, BK] + + if V > 64: + p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dk += tl.dot(b_v, tl.trans(b_dh2).to(b_v.dtype)) + + if V > 128: + p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dk += tl.dot(b_v, tl.trans(b_dh3).to(b_v.dtype)) + + if V > 192: + p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dk += tl.dot(b_v, tl.trans(b_dh4).to(b_v.dtype)) + + b_dk += tl.load(p_dk, boundary_check=(0, 1)) + + tl.store(p_dk2, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + # Update dh, 按照K切分，收集所有V维度，q一次就好，wdo要收集所有 + + p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) # [BK, BT] + b_q = tl.load(p_q, boundary_check=(0, 1)) + + if V > 0: + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] + b_w = tl.load(p_w, boundary_check=(0, 1)) + p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + if USE_GV: + o_v1 = tl.arange(0, 64) + b_gv_last1 = tl.load(gv + last_idx * H*V + o_v1, mask=(o_v1 < V), other=0.) + b_dh1 *= exp(b_gv_last1[None, :]) + b_do *= exp(b_gv) + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - \ + tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) # [BK, BT] @ [BT, BV] - [BK, BT] @ [BT, BV] + + if V > 64: + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] + b_w = tl.load(p_w, boundary_check=(0, 1)) + p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + if USE_GV: + o_v2 = 64 + o_v1 + b_gv_last2 = tl.load(gv + last_idx * H*V + o_v2, mask=(o_v2 < V), other=0.) + b_dh2 *= exp(b_gv_last2[None, :]) + b_do *= exp(b_gv) + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) + + if V > 128: + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] + b_w = tl.load(p_w, boundary_check=(0, 1)) + p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + if USE_GV: + o_v3 = 128 + o_v1 + b_gv_last3 = tl.load(gv + last_idx * H*V + o_v3, mask=(o_v3 < V), other=0.) + b_dh3 *= exp(b_gv_last3[None, :]) + b_do *= exp(b_gv) + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) + + if V > 192: + p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] + b_w = tl.load(p_w, boundary_check=(0, 1)) + p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + if USE_GV: + o_v4 = 192 + o_v1 + b_gv_last4 = tl.load(gv + last_idx * H*V + o_v4, mask=(o_v4 < V), other=0.) + b_dh4 *= exp(b_gv_last4[None, :]) + b_do *= exp(b_gv) + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + if V > 64: + p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + if V > 128: + p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + if V > 192: + p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_bwd_dhu( + q: torch.Tensor, + vg: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dk: torch.Tensor, + gv: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + states_in_fp32: bool = False +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *q.shape, do.shape[-1] + BT = 64 + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + dh = q.new_empty(B, NT, H, K, V, dtype=q.dtype if not states_in_fp32 else torch.float) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dk2 = torch.empty_like(dk) + + def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) + chunk_oja_bwd_kernel_dhu_blockdim64[grid]( + q=q, + vg=vg, + w=w, + gv=gv, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dk=dk, + dk2=dk2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dh, dh0, dk2 + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['BT'] +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_dqkvg( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_i = tl.arange(0, BT) + o_t = min(i_t * BT + BT, T) + m_s = o_i[:, None] >= o_i[None, :] + + p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + m_v = o_v < V + + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_gv = exp(b_gn[None, :] - b_g) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g) * scale).to(b_do.dtype) + # [BK, BV] + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BV] + b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp(b_gn) + + b_dh = b_dh.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) + b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) + # [BT, BV] + b_dv = tl.dot(b_k, b_dh) * b_gv + # [BV] + b_dg += tl.sum(b_dv * b_v, 0) + + if i_k == 0: + b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] + else: + b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] + + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BT] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BT, BK] + b_dq += tl.dot(b_dA, b_k) + b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_GV': lambda args: args['gv'] is not None, + 'HAVE_GK': lambda args: args['dgk'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'USE_GV'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_bwd_kernel_dvwg_h( + k, + v, + gv, + h, + dh, + dk, + dw, + dv, + dgv_last, + dgk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_GV: tl.constexpr, + HAVE_GK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + gv += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K*V + dh += (i_tg * H + i_h).to(tl.int64) * K*V + dk += (bos * H + i_h) * K + dw += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + dgv_last += (bos * H + i_h) * V + + b_dvg = tl.zeros([BT, BV], dtype=tl.float32) + b_dw = tl.zeros([BT, BV], dtype=tl.float32) + b_dgv_last = tl.zeros([BV,], dtype=tl.float32) + + if USE_GV: + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + p_gn = gv + (min(T, i_t * BT + BT) - 1) * H*V + o_v + p_gv = tl.make_block_ptr(gv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_gn = tl.load(p_gn, mask=m_v, other=0) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) # BT BK + b_dk = tl.load(p_dk, boundary_check=(0, 1)) # BT BK + b_h = tl.load(p_h, boundary_check=(0, 1)) # BK BV + b_dh = tl.load(p_dh, boundary_check=(0, 1)) # BK BV + + b_dvg += tl.dot(b_k, b_dh.to(b_k.dtype)) # BT BK @ BK BV -> BT BV + b_dw += tl.dot(b_dk.to(b_k.dtype), b_h.to(b_k.dtype)) # BT BK @ BK BV -> BT BV + b_dgv_last += tl.sum((b_h * b_dh) * exp(b_gn), axis=0) + + if USE_GV: + b_dv = b_dvg * exp(b_gn[None, :] - b_gv) + + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dw = tl.make_block_ptr(dw, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dgv_last = tl.make_block_ptr(dgv_last, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + b_dgv_last += tl.sum(b_dv * b_v, axis=0) + + # 留给GSA2的接口 + if HAVE_GK: + dgk += (bos * H + i_h) * V + p_dgk = tl.make_block_ptr(dgk, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) + b_dgv_last = b_dgk + b_dgv_last[None, :] + else: + b_dgv_last = tl.zeros([BT, BV], dtype=tl.float32) + b_dgv_last[None, :] + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dgv_last, b_dgv_last.to(p_dgv_last.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_bwd_dvwg_h( + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + dk: torch.Tensor, + gv: torch.Tensor | None = None, + dgk: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + B, T, H, K, V = *k.shape, v.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NV = triton.cdiv(V, BV) + dv = torch.empty_like(v, dtype=torch.float) + dw = torch.empty_like(v) + dgv_last = torch.empty_like(gv) + + grid = (NV, NT, B * H) + chunk_oja_bwd_kernel_dvwg_h[grid]( + k=k, + v=v, + gv=gv, + h=h, + dh=dh, + dw=dw, + dk=dk, + dv=dv, + dgv_last=dgv_last, + dgk=dgk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dv, dw, dgv_last diff --git a/fla/ops/gated_oja_rule/chunk_kkt.py b/fla/ops/gated_oja_rule/chunk_kkt.py new file mode 100644 index 0000000000000000000000000000000000000000..5753ac6612d660c98d9fc21faf0e4d567574ea95 --- /dev/null +++ b/fla/ops/gated_oja_rule/chunk_kkt.py @@ -0,0 +1,521 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'BT', 'IS_VARLEN'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A *= exp(b_g_diff) + b_A *= b_b[:, None] + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"] +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H*K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) * exp(b_gn[:, None] - b_gk) + # [BC, BC] + b_A += tl.dot(b_k, b_kt) + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + ], + key=["BK", "BT"] +) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H*BT + i_h * BT + i_i * BC + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + + b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_b, mask=m_A, other=0)[:, None] + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_kt = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_A = tl.sum(b_k * b_kt[None, :] * exp(b_g - b_gk[None, :]), 1) + b_A = tl.where(o_i > j, b_A, 0.) + + tl.store(A + o_A + j, b_A, mask=m_A) + p_kt += H*K + p_gk += H*K + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BK', 'NC', 'BT'], +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_scaled_dot_kkt_bwd_kernel_gk( + k, + g, + beta, + dA, + dk, + dg, + db, + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = i_c // NC, i_c % NC + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + if i_t * BT + i_i * BC >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + beta += bos * H + i_h + + dA += (bos * H + i_h) * BT + dk += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + db += (i_k * all + bos) * H + i_h + + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + # [BC] + b_b = tl.load(p_b, boundary_check=(0,)) + if i_i > 0: + p_gn = g + (i_t * BT + i_i * BC) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[None, :] - b_gk) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + b_dkb = tl.dot(b_dA, b_kg) * exp(b_g - b_gn[None, :]) + b_dk += b_dkb + + o_i = tl.arange(0, BC) + m_dA = (i_t * BT + i_i * BC + o_i) < T + o_dA = (i_t * BT + i_i * BC + o_i) * H*BT + i_i * BC + p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k + p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k + + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) + # [BK] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_dkb = tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) + b_dk += b_dkb + + p_kj += H*K + p_gkj += H*K + b_db = tl.sum(b_dk * b_k, 1) + b_dk *= b_b[:, None] + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + tl.debug_barrier() + # [BC, BK] + b_dkt = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = g + (min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(i_i + 1, NC): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC] + b_b = tl.load(p_b, boundary_check=(0,)) + # [BC, BK] + b_kb = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) * b_b[:, None] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dkt += tl.dot(b_dA, b_kbg) + b_dkt *= exp(b_gn[None, :] - b_g) + o_dA = (i_t * BT + i_i * BC) * H*BT + i_i * BC + o_i + p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k + p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k + p_bj = beta + (i_t * BT + i_i * BC) * H + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dA = tl.load(dA + o_dA + j * H*BT) + # [BK,] + b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + b_kbgj = b_kbj[None, :] * exp(b_gkj[None, :] - b_g) + # [BC, BK] + m_i = o_i[:, None] <= j + b_dkt += tl.where(m_i, b_dA[:, None] * b_kbgj, 0.) + + p_kj += H*K + p_gkj += H*K + p_bj += H + b_dg = (b_dk - b_dkt) * b_k + b_dk += b_dkt + + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_kkt_fwd( + k: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_indices (torch.LongTensor): + Pre-computed chunk indices. Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + if gk is None: + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + ) + return A + + BC = min(16, BT) + NC = triton.cdiv(BT, BC) + BK = max(triton.next_power_of_2(K), 16) + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + k=k, + g=gk, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + k=k, + g=gk, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + return A + + +def chunk_scaled_dot_kkt_bwd_gk( + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dA: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64 +): + B, T, H, K = k.shape + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dk = torch.empty_like(k, dtype=torch.float) + dg = torch.empty_like(g, dtype=torch.float) + db = beta.new_empty(NK, *beta.shape, dtype=torch.float) + grid = (NK, NT * NC, B * H) + chunk_scaled_dot_kkt_bwd_kernel_gk[grid]( + k=k, + g=g, + beta=beta, + dA=dA, + dk=dk, + dg=dg, + db=db, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + db = db.sum(0) + + return dk, dg, db diff --git a/fla/ops/gated_oja_rule/chunk_o.py b/fla/ops/gated_oja_rule/chunk_o.py new file mode 100644 index 0000000000000000000000000000000000000000..2f02332c95e5fd81c30ac02b4fc4bfffa7b9b7fb --- /dev/null +++ b/fla/ops/gated_oja_rule/chunk_o.py @@ -0,0 +1,674 @@ + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import check_shared_mem, is_nvidia_hopper + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'] +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_fwd_inter( + q, + k, + h, + gv, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BT] + b_A += tl.dot(b_q, b_k) + p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o = b_o * exp(b_g) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + # [BT, BT] + b_A = tl.where(m_s, b_A, 0.) + if i_v == 0: + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_fwd_intra( + v, + gv, + o, + A, + cu_seqlens, + chunk_indices, + T, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = gv + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_o = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(0, i_i): + p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[None, :] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_o += tl.dot(b_A, b_vg) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o *= exp(b_g - b_gn[None, :]) + + o_i = tl.arange(0, BC) + o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_gv = gv + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + # [BC,] + b_A = tl.load(A + o_A + j, mask=m_A, other=0) + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC, BV] + b_vg = b_v[None, :] * exp(b_g - b_gv[None, :]) + # avoid 0 * inf = inf + b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) + p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor, + h: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64 +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + HQ = q.shape[2] + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NG = HQ // H + + o = v.new_empty(B, T, HQ, V) + A = q.new_empty(B, T, HQ, BT) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) + chunk_oja_fwd_inter[grid]( + q, + k, + h, + gv, + o, + A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + ) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + chunk_oja_fwd_intra[grid]( + v, + gv, + o, + A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + num_warps=4, + num_stages=2 + ) + return A, o + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8] + ], + key=["BT"] +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_bwd_kernel_dA( + v, + gv, + do, + dA, + chunk_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + # [BC, BC] + b_dA = tl.zeros([BC, BC], dtype=tl.float32) + if i_i > i_j: + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gn = gv + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + p_g = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0.) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g - b_gn[None, :]) * scale).to(b_do.dtype) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[:, None] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_dA = tl.dot(b_do, b_vg) + elif i_i == i_j: + p_g = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + p_gv = gv + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + m_v = o_v < V + + o_i = tl.arange(0, BC) + # [BC, BC] + m_dA = o_i[:, None] >= o_i[None, :] + for j in range(0, min(BC, T - i_t * BT - i_j * BC)): + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC,] + b_dAj = tl.sum(b_do * b_v[None, :] * exp(b_g - b_gv[None, :]), 1) + b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) + + p_v += H*V + p_gv += H*V + b_dA = tl.where(m_dA, b_dA, 0.) + + p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_bwd_dA( + v: torch.Tensor, + gv: torch.Tensor, + do: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64 +): + B, T, H, V = v.shape + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NV = triton.cdiv(V, BV) + + dA = v.new_empty(NV, B, T, H, BT) + # 计算dA + grid = (NV, NT * NC * NC, B * H) + chunk_oja_bwd_kernel_dA[grid]( + v, + gv, + do, + dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + ) + dA = dA.sum(0, dtype=dA.dtype) + + return dA + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['BT'] +) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_bwd_kernel_dqk( + q, + k, + h, + gv, + A, + dq, + dk, + dA, + do, + scale, + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + # [B, T, H, BT] + p_q = tl.make_block_ptr(q + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + ((i_k*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + + # 先计算do对应的dq + for i_v in range(tl.cdiv(V, BV)): + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = tl.make_block_ptr(gv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_do = (b_do * exp(b_gv) * scale).to(b_do.dtype) + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + + # 接着计算dA对应的dq, dk + p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BT] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BT, BK] + b_dq += tl.dot(b_dA.to(b_q.dtype), b_k) + b_dk = tl.dot(tl.trans(b_dA).to(b_q.dtype), b_q) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_bwd_dqk( + q: torch.Tensor, + k: torch.Tensor, + h: torch.Tensor, + gv: torch.Tensor, + dA: torch.Tensor, + do: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64 +): + B, T, H, K, V = *q.shape, gv.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NK = triton.cdiv(K, BK) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + A = dA.new_empty(NK, B, T, H, BT) + # 计算dA + grid = (NK, NT, B * H) + chunk_oja_bwd_kernel_dqk[grid]( + q, + k, + h, + gv, + A, + dq, + dk, + dA, + do, + scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV + ) + + A = A.sum(0, dtype=A.dtype) + + return A, dq, dk + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def chunk_oja_bwd_kernel_dv_o( + v, + g, + o, + A, + do, + dv, + dv2, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_dvg = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos*H+i_h) * BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + # [BC, BV] + b_dvg += tl.dot(b_A, b_do.to(b_A.dtype)) + b_dv = b_dvg * exp(b_gn[None, :] - b_gv) + + o_i = tl.arange(0, BC) + o_c = i_i * BC + tl.arange(0, BC) + + p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v + p_A = A + (bos + i_t*BT + i_i*BC) * H*BT + i_h * BT + o_c + p_do = do + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_A = tl.load(p_A) + # [BV,] + b_g = tl.load(p_g, mask=m_v, other=0) + b_do = tl.load(p_do, mask=m_v, other=0) + # [BC, BV] + m_i = o_i[:, None] <= j + b_dv += tl.where(m_i, exp(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) + + p_g += H * V + p_A += H * BT + p_do += H * V + p_o = tl.make_block_ptr(o + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2 + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + + b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) + b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) + b_dg = b_o * b_do - b_v * b_dv + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_oja_bwd_dv_o( + v: torch.Tensor, + gv: torch.Tensor, + o: torch.Tensor, + A: torch.Tensor, + dv: torch.Tensor, + do: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64 +): + B, T, H, V = v.shape + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + dv2 = torch.empty_like(v, dtype=torch.float) + dgv = torch.empty_like(gv) + # 计算dA + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * H) + chunk_oja_bwd_kernel_dv_o[grid]( + v=v, + g=gv, + o=o, + A=A, + do=do, + dv=dv, + dv2=dv2, + dg=dgv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + num_warps=4, + num_stages=2 + ) + return dv2, dgv diff --git a/fla/ops/gated_oja_rule/fused_recurrent.py b/fla/ops/gated_oja_rule/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..b40aced3ac97f786a051d0971708079ca9dd7a14 --- /dev/null +++ b/fla/ops/gated_oja_rule/fused_recurrent.py @@ -0,0 +1,263 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_oja_fwd_kernel( + q, + k, + v, + gv, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_GV: tl.constexpr, + USE_Q_L2NORM: tl.constexpr, + USE_K_L2NORM: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if USE_GV: + p_gv = gv + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + bos * HV + i_hv + else: + p_beta = beta + (bos * HV + i_hv) * V + o_v + + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + if USE_Q_L2NORM: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + if USE_K_L2NORM: + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta).to(tl.float32) + else: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + + # [BK, BV] + if USE_GV: + b_gv = tl.load(p_gv, mask=mask_v, other=0).to(tl.float32) + b_h *= exp(b_gv[None, :]) + + b_k = b_beta * (b_k - tl.sum(b_h * b_v[None, :], 1)) + b_h += b_k[:, None] * b_v + + # [BV] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_v += HV*V + if USE_GV: + p_gv += HV*V + p_beta += HV * (1 if IS_BETA_HEADWISE else V) + p_o += HV*V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_oja_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_q_l2norm: bool = False, + use_k_l2norm: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + assert V <= 128 + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 256) + NV = triton.cdiv(V, BV) + num_stages = 3 + num_warps = 1 + + o = torch.empty_like(v) + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + grid = (NV, N * HV) + fused_recurrent_oja_fwd_kernel[grid]( + q=q, + k=k, + v=v, + gv=gv, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim != v.ndim, + USE_Q_L2NORM=use_q_l2norm, + USE_K_L2NORM=use_k_l2norm, + num_warps=num_warps, + num_stages=num_stages, + ) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_q_l2norm: bool = False, + use_k_l2norm: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, final_state = fused_recurrent_oja_fwd( + q=q, + k=k, + v=v, + gv=gv, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_q_l2norm=use_q_l2norm, + use_k_l2norm=use_k_l2norm, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps." + ) + + +def fused_recurrent_gated_oja_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_q_l2norm: bool = False, + use_k_l2norm: bool = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + + if 'use_qk_l2norm_in_kernel' in kwargs and (not use_q_l2norm and not use_k_l2norm): + use_q_l2norm = True + use_k_l2norm = True + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if beta is None: + beta = torch.ones_like(q[..., 0]) + + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + gv, + beta, + scale, + initial_state, + output_final_state, + use_q_l2norm, + use_k_l2norm, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/gated_oja_rule/wy_fast.py b/fla/ops/gated_oja_rule/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..7ff94e5c641c186bc7ceccb78552f0c5768a848a --- /dev/null +++ b/fla/ops/gated_oja_rule/wy_fast.py @@ -0,0 +1,290 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import check_shared_mem + + +@triton.heuristics({ + 'STORE_VG': lambda args: args['vg'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + vg, + beta, + w, + u, + A, + gv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_VG: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_w = tl.make_block_ptr(w + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = b_v * b_b[:, None] + + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vb *= exp(b_gv) + if STORE_VG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + b_gn = tl.load(gv + ((bos + last_idx) * H + i_h) * V + o_v, mask=m_v, other=0.) + b_vg = b_v * exp(b_gn - b_gv) + + p_vg = tl.make_block_ptr(vg + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_vg, b_vg.to(p_vg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_vb.to(b_v.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_b[:, None]).to(b_k.dtype) + b_u = tl.dot(b_A, b_kb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'] +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + gv, + A, + dA, + dw, + du, + dk, + dv, + db, + dgv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_b = tl.load(p_b, boundary_check=(0,)) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_gv_exp = exp(tl.load(p_gv, boundary_check=(0, 1))) + b_vbg = b_v * b_b[:, None] * b_gv_exp + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dw, tl.trans(b_vbg).to(b_dw.dtype)) + b_dvbg = tl.dot(b_A, b_dw) + b_dv = b_dvbg * b_gv_exp * b_b[:, None] + b_db += tl.sum(b_dvbg * b_v * b_gv_exp, 1) + b_dgv = b_dvbg * b_vbg + + p_dgv = tl.make_block_ptr(dgv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_b[:, None]).to(b_k.dtype) # BT BK + b_du = tl.load(p_du, boundary_check=(0, 1)) # BT BK + b_dA += tl.dot(b_du, tl.trans(b_kb)) # BT BT + b_dkb = tl.dot(b_A, b_du) # BT BK + b_dk = b_dkb * b_b[:, None] + b_db += tl.sum(b_dkb * b_k, 1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + b_dA = tl.where(m_A, -b_dA, 0) + + # if USE_GV: + p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + gv: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(v) + u = torch.empty_like(k) + vg = torch.empty_like(v) if gv is not None else None + recompute_w_u_fwd_kernel[(NT, B*H)]( + k=k, + v=v, + vg=vg, + beta=beta, + w=w, + u=u, + A=A, + gv=gv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u, vg + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + gv: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = torch.empty_like(k) + dv = torch.empty_like(v, dtype=torch.float) + + dgv = torch.empty_like(gv, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + gv=gv, + A=A, + dA=dA, + dw=dw, + du=du, + dk=dk, + dv=dv, + db=db, + dgv=dgv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + + return dk, dv, db, dgv, dA diff --git a/fla/ops/generalized_delta_rule/README.md b/fla/ops/generalized_delta_rule/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f96c22f44a51ad3e6fdeb824eb2aded660223600 --- /dev/null +++ b/fla/ops/generalized_delta_rule/README.md @@ -0,0 +1,37 @@ +# Generalized Delta Rule + +In delta rule we have the recurrence: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I}-\beta_t \mathbf{k}_t\mathbf{k}_t^T) + \beta_t \mathbf{v}_t\mathbf{k}_t^T +``` + +This repository implements a delta rule variant where $\mathbf{I}$ is not necessarily an identity matrix; $\mathbf{k}_t$ in $\mathbf{I} - \beta_t \mathbf{k}_t\mathbf{k}_t^T$ might be different from input $\mathbf{k}_t$ in $\mathbf{v}_t\mathbf{k}_t^T$. + +## IPLR (Identity Plus Low Rank) + +The first variant is IPLR, where we have: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I}+\mathbf{a}_t\mathbf{b}_t^T) + \mathbf{v}_t\mathbf{k}_t^T +``` + +When $\mathbf{a}_t = -\beta_t \mathbf{k}_t$, $\mathbf{b}_t = \mathbf{k}_t$, $\mathbf{v}_t= \beta_t \mathbf{v}_t$, we recover the original delta rule. Since here the transition matrix is identity-plus-low-rank, we refer to this variant as IPLR. + +### Numerical Stability + +$\mathbf{a}_t$ and $\mathbf{b}_t$ must be in opposite directions, that is, $\mathbf{b}_t = \lambda_t \mathbf{a}_t$ where $\lambda_t < 0$. For an understanding of why this is necessary, you can derive the eigenvalues of the transition matrix. + +## DPLR (Diagonal Plus Low Rank) + +The second variant is DPLR, where we have: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{D}_t+\mathbf{a}_t\mathbf{b}_t^T) + \mathbf{v}_t\mathbf{k}_t^T +``` + +Here, $\mathbf{I}$ is replaced by a diagonal matrix $\mathbf{D}_t$. This transition matrix structure has been utilized in RWKV7. + +## Efficient Chunkwise Implementation + +For detailed information about efficient chunkwise implementation, please refer to our [technical note](https://drive.google.com/file/d/1rJbO3dU4fe7OKG3w7Yg058z_BNIuavNF/view?usp=sharing). diff --git a/fla/ops/generalized_delta_rule/__init__.py b/fla/ops/generalized_delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..33875bc69eae5bea4ec131e11b773c69b6e1c093 --- /dev/null +++ b/fla/ops/generalized_delta_rule/__init__.py @@ -0,0 +1,9 @@ +from .dplr import chunk_dplr_delta_rule, fused_recurrent_dplr_delta_rule +from .iplr import chunk_iplr_delta_rule, fused_recurrent_iplr_delta_rule + +__all__ = [ + 'chunk_dplr_delta_rule', + 'fused_recurrent_dplr_delta_rule', + 'chunk_iplr_delta_rule', + 'fused_recurrent_iplr_delta_rule', +] diff --git a/fla/ops/generalized_delta_rule/dplr/__init__.py b/fla/ops/generalized_delta_rule/dplr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..95b9d07a31be8868d25c7e86d564b399eb5b1532 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/__init__.py @@ -0,0 +1,7 @@ +from .chunk import chunk_dplr_delta_rule +from .fused_recurrent import fused_recurrent_dplr_delta_rule + +__all__ = [ + 'chunk_dplr_delta_rule', + 'fused_recurrent_dplr_delta_rule', +] diff --git a/fla/ops/generalized_delta_rule/dplr/chunk.py b/fla/ops/generalized_delta_rule/dplr/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..1797ea74ab37527ff26e94d5129bb96c39e18982 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk.py @@ -0,0 +1,434 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.generalized_delta_rule.dplr.chunk_A_bwd import chunk_dplr_bwd_dqk_intra +from fla.ops.generalized_delta_rule.dplr.chunk_A_fwd import chunk_dplr_fwd_intra +from fla.ops.generalized_delta_rule.dplr.chunk_h_bwd import chunk_dplr_bwd_dhu +from fla.ops.generalized_delta_rule.dplr.chunk_h_fwd import chunk_dplr_fwd_h +from fla.ops.generalized_delta_rule.dplr.chunk_o_bwd import chunk_dplr_bwd_dAu, chunk_dplr_bwd_dv, chunk_dplr_bwd_o +from fla.ops.generalized_delta_rule.dplr.chunk_o_fwd import chunk_dplr_fwd_o +from fla.ops.generalized_delta_rule.dplr.wy_fast_bwd import chunk_dplr_bwd_wy +from fla.ops.generalized_delta_rule.dplr.wy_fast_fwd import prepare_wy_repr_fwd +from fla.ops.rwkv6.chunk import chunk_rwkv6_fwd_cumsum +from fla.ops.utils import prepare_chunk_indices +from fla.utils import TRITON_ABOVE_3_4_0, autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_dplr_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + safe_gate: bool = False, + chunk_indices: torch.LongTensor | None = None, + disable_recompute: bool = False, +): + gi, ge = chunk_rwkv6_fwd_cumsum(gk, chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + + A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_dplr_fwd_intra( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + scale=scale, + cu_seqlens=cu_seqlens, + safe_gate=safe_gate, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # A_ab, A_ak, gi, ge torch.float32 + # A_qk, A_qb, qg, kg, ag, bg, dtype=q.dtype, eg: bf16 + w, u, A_ab_inv = prepare_wy_repr_fwd( + ag=ag, + A_ab=A_ab, + A_ak=A_ak, + v=v, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + h, v_new, final_state = chunk_dplr_fwd_h( + kg=kg, + bg=bg, + v=v, + w=w, + u=u, + gk=gi, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + o = chunk_dplr_fwd_o( + qg=qg, + v=v, + v_new=v_new, + A_qk=A_qk, + A_qb=A_qb, + h=h, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + if disable_recompute: + return o, final_state, (gi, ge, A_qk, A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv) + else: + return o, final_state, None + + +class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + safe_gate: bool = False, + chunk_size: int | None = None, + disable_recompute: bool = False, + ): + # Due to gate numerical stability consideration, we only support chunk_size=16 when safe_gate=True + # And in practice, chunk_size=16 is sufficient for no safe gate situations. + # It's different from the other chunk implementations. + if chunk_size is None: + chunk_size = 16 + elif TRITON_ABOVE_3_4_0: + chunk_size = chunk_size + else: + # Avoid Triton Compiler error + warnings.warn( + "Set chunk_size to 16, to avoid triton compiler erorr. " + f"original chunk_size {chunk_size}", + category=RuntimeWarning, + stacklevel=2, + ) + chunk_size = 16 + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + o, final_state, cache = chunk_dplr_fwd( + q=q, + k=k, + v=v, + a=a, + b=b, + gk=gk, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + safe_gate=safe_gate, + chunk_indices=chunk_indices, + disable_recompute=disable_recompute, + ) + + if disable_recompute: + gi, ge, A_qk, A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv = cache + ctx.save_for_backward(q, k, v, a, b, gk, initial_state, gi, ge, A_qk, + A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv) + else: + ctx.save_for_backward(q, k, v, a, b, gk, initial_state) + ctx.cu_seqlens = cu_seqlens + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.chunk_indices = chunk_indices + ctx.safe_gate = safe_gate + ctx.disable_recompute = disable_recompute + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + if ctx.disable_recompute: + ( + q, k, v, a, b, gk, initial_state, + gi, ge, A_qk, A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv, + ) = ctx.saved_tensors + else: + q, k, v, a, b, gk, initial_state = ctx.saved_tensors + chunk_size = ctx.chunk_size + cu_seqlens = ctx.cu_seqlens + scale = ctx.scale + + if not ctx.disable_recompute: + # ******* start recomputing everything, otherwise i believe the gpu memory will be exhausted ******* + gi, ge = chunk_rwkv6_fwd_cumsum(gk, chunk_size, cu_seqlens=cu_seqlens, chunk_indices=ctx.chunk_indices) + + A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_dplr_fwd_intra( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + scale=scale, + cu_seqlens=cu_seqlens, + safe_gate=ctx.safe_gate, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + w, u, A_ab_inv = prepare_wy_repr_fwd( + ag=ag, + A_ab=A_ab, + A_ak=A_ak, + v=v, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + del A_ab + h, v_new, _ = chunk_dplr_fwd_h( + kg=kg, + bg=bg, + v=v, + w=w, + u=u, + gk=gi, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + del u + # ******* end of recomputation ******* + # A_ak, A_ab_inv, gi, ge torch.float32 + # A_qk, A_qb, qg, kg, ag, bg, v_new dtype=q.dtype, eg: bf16 + + dv_new_intra, dA_qk, dA_qb = chunk_dplr_bwd_dAu( + v=v, + v_new=v_new, + do=do, + A_qb=A_qb, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + + dh, dh0, dv_new = chunk_dplr_bwd_dhu( + qg=qg, + bg=bg, + w=w, + gk=gi, + h0=initial_state, + dht=dht, + do=do, + dv=dv_new_intra, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + + dv = chunk_dplr_bwd_dv( + A_qk=A_qk, + kg=kg, + do=do, + dh=dh, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + del A_qk + + dqg, dkg, dw, dbg, dgk_last = chunk_dplr_bwd_o( + k=kg, + b=bg, + v=v, + v_new=v_new, + do=do, + h=h, + dh=dh, + dv=dv_new, + w=w, + gk=gi, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + scale=scale, + chunk_indices=ctx.chunk_indices, + ) + del v_new + + dA_ab, dA_ak, dv, dag = chunk_dplr_bwd_wy( + A_ab_inv=A_ab_inv, + A_ak=A_ak, + v=v, + ag=ag, + dw=dw, + du=dv_new, + dv0=dv, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=ctx.chunk_indices, + ) + del A_ak + + dq, dk, da, db, dgk = chunk_dplr_bwd_dqk_intra( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + dAqk=dA_qk, + dAqb=dA_qb, + dAak=dA_ak, + dAab=dA_ab, + dgk_last=dgk_last, + dqg=dqg, + dkg=dkg, + dag=dag, + dbg=dbg, + chunk_size=chunk_size, + scale=scale, + cu_seqlens=cu_seqlens, + safe_gate=ctx.safe_gate, + chunk_indices=ctx.chunk_indices, + ) + + return dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), dgk.to(gk), None, dh0, None, None, None, None, None, None + + +@torch.compiler.disable +def chunk_dplr_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, + safe_gate: bool = False, + chunk_size: int | None = None, + disable_recompute: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + a (torch.Tensor): + activations of shape `[B, T, H, K]`. + b (torch.Tensor): + betas of shape `[B, T, H, K]`. + gk (torch.Tensor): + gk of shape `[B, T, H, K]`. decay term in log space! + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cu_seqlens_cpu (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + safe_gate (bool): + Whether the kernel can assume the input gate values `g` are in a safe range. + When `True`, the kernel can use M=16 TensorCore acceleration. + The safe range is approximately [-5, 0). Default: `False`. + chunk_size (Optional[int]): + Chunk size for the chunked computation. Default: `None`, which means 16. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if q.dtype == torch.float32: + warnings.warn( + """ChunkDeltaRuleFunction does not support float32 on some platforms. Please use bfloat16/float16. + If you want to use float32, please solve the issue by yourself.""", + category=RuntimeWarning, + stacklevel=2, + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + scale = k.shape[-1] ** -0.5 if scale is None else scale + o, final_state = ChunkDPLRDeltaRuleFunction.apply( + q, + k, + v, + a, + b, + gk, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + safe_gate, + chunk_size, + disable_recompute, + ) + return o, final_state diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..4336224d565a98cf0d5d6595ad65ea75a1fda524 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py @@ -0,0 +1,542 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, gather +from fla.utils import IS_AMD, IS_GATHER_SUPPORTED, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BK', 'BT', 'K'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_bwd_kernel_intra( + q, + k, + a, + b, + gi, + ge, + dAqk, + dAqb, + dAak, + dAab, + dq, + dk, + da, + db, + dqg, + dkg, + dag, + dbg, + dgk, + dgk_offset, + cu_seqlens, + chunk_indices, + scale: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + GATHER_SUPPORTED: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + + if i_t * BT >= T: + return + + # offset calculation + ge += (bos*H + i_h) * K + gi += (bos*H + i_h) * K + q += (bos*H + i_h) * K + a += (bos*H + i_h) * K + b += (bos*H + i_h) * K + k += (bos*H + i_h) * K + dq += (bos*H + i_h) * K + dk += (bos*H + i_h) * K + da += (bos*H + i_h) * K + db += (bos*H + i_h) * K + dqg += (bos*H + i_h) * K + dag += (bos*H + i_h) * K + dkg += (bos*H + i_h) * K + dbg += (bos*H + i_h) * K + dgk += (bos*H + i_h) * K + dgk_offset += (bos*H + i_h) * K + dAqk += (bos*H + i_h) * BT + dAqb += (bos*H + i_h) * BT + dAak += (bos*H + i_h) * BT + dAab += (bos*H + i_h) * BT + + stride_qk = H*K + stride_A = H*BT + + p_ge = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_gi = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_ge = tl.load(p_ge, boundary_check=(0, 1)) + b_gi = tl.load(p_gi, boundary_check=(0, 1)) + b_dq = tl.zeros([BC, BK], dtype=tl.float32) + b_da = tl.zeros([BC, BK], dtype=tl.float32) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + b_db = tl.zeros([BC, BK], dtype=tl.float32) + # intra chunk gradient calculation + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) + p_dAab = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) + p_dAqb = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) + p_dAak = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) + o_i = tl.arange(0, BC) + p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) + p_a = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) + b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) + b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) + + # inter chunk gradient calculation + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # intra chunk gradient calculation + for j in range(0, min(BC, T - i_t * BT)): + # trick to index the block + if GATHER_SUPPORTED: + row_idx = tl.full([1, BK], j, dtype=tl.int16) + col_idx = tl.full([BC, 1], j, dtype=tl.int16) + row_idx_bc = tl.full([1, BC], j, dtype=tl.int16) + # [1, BK] + b_kj = gather(b_k, row_idx, axis=0) + b_bj = gather(b_b, row_idx, axis=0) + b_gij = gather(b_gi, row_idx, axis=0) + b_gej = gather(b_ge, row_idx, axis=0) + b_qj = gather(b_q, row_idx, axis=0) + b_aj = gather(b_a, row_idx, axis=0) + # [BC, 1] + b_dAqk_j = gather(b_dAqk, col_idx, axis=1) + b_dAab_j = gather(b_dAab, col_idx, axis=1) + b_dAqb_j = gather(b_dAqb, col_idx, axis=1) + b_dAak_j = gather(b_dAak, col_idx, axis=1) + # [1, BC] -> [BC, 1] + b_dA_qk_j = tl.sum(gather(b_dAqk, row_idx_bc, axis=0), 0)[:, None] + b_dA_qk_j = tl.sum(gather(b_dAqk, row_idx_bc, axis=0), 0)[:, None] + b_dA_ab_j = tl.sum(gather(b_dAab, row_idx_bc, axis=0), 0)[:, None] + b_dA_qb_j = tl.sum(gather(b_dAqb, row_idx_bc, axis=0), 0)[:, None] + b_dA_ak_j = tl.sum(gather(b_dAak, row_idx_bc, axis=0), 0)[:, None] + else: + mask_idx = tl.arange(0, BC) == j + b_kj = tl.sum(tl.where(mask_idx[:, None], b_k, 0), 0)[None, :] + b_bj = tl.sum(tl.where(mask_idx[:, None], b_b, 0), 0)[None, :] + b_gij = tl.sum(tl.where(mask_idx[:, None], b_gi, 0), 0)[None, :] + b_gej = tl.sum(tl.where(mask_idx[:, None], b_ge, 0), 0)[None, :] + b_dAqk_j = tl.sum(tl.where(mask_idx[None, :], b_dAqk, 0), 1)[:, None] + b_dAab_j = tl.sum(tl.where(mask_idx[None, :], b_dAab, 0), 1)[:, None] + b_dAqb_j = tl.sum(tl.where(mask_idx[None, :], b_dAqb, 0), 1)[:, None] + b_dAak_j = tl.sum(tl.where(mask_idx[None, :], b_dAak, 0), 1)[:, None] + b_dA_qk_j = tl.sum(tl.where(mask_idx[:, None], b_dAqk, 0), 0)[:, None] + b_dA_ab_j = tl.sum(tl.where(mask_idx[:, None], b_dAab, 0), 0)[:, None] + b_dA_qb_j = tl.sum(tl.where(mask_idx[:, None], b_dAqb, 0), 0)[:, None] + b_dA_ak_j = tl.sum(tl.where(mask_idx[:, None], b_dAak, 0), 0)[:, None] + # [1, BK] b_qj, b_aj + b_qj = tl.sum(tl.where(mask_idx[:, None], b_q, 0), 0)[None, :] + b_aj = tl.sum(tl.where(mask_idx[:, None], b_a, 0), 0)[None, :] + + m_e = o_i[:, None] > j + m_i = o_i[:, None] >= j + tmp1 = exp(b_gi - b_gij) + tmp2 = exp(b_ge - b_gij) + b_dq += tl.where(m_i, b_dAqk_j * b_kj * tmp1, 0.) + b_dq += tl.where(m_i, b_dAqb_j * b_bj * tmp1, 0.) + b_da += tl.where(m_e, b_dAab_j * b_bj * tmp2, 0.) + b_da += tl.where(m_e, b_dAak_j * b_kj * tmp2, 0.) + + m_i = o_i[:, None] <= j + m_e = o_i[:, None] < j + tmp1 = exp(b_gij - b_gi) + tmp2 = exp(b_gej - b_gi) + b_dk += tl.where(m_i, b_dA_qk_j * b_qj * tmp1, 0.) + b_dk += tl.where(m_e, b_dA_ak_j * b_aj * tmp2, 0.) + b_db += tl.where(m_i, b_dA_qb_j * b_qj * tmp1, 0.) + b_db += tl.where(m_e, b_dA_ab_j * b_aj * tmp2, 0.) + + # post processing + p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_da = tl.make_block_ptr(da, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dqg = tl.make_block_ptr(dqg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dkg = tl.make_block_ptr(dkg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dag = tl.make_block_ptr(dag, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dbg = tl.make_block_ptr(dbg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_gn = gi + (min(i_t * BT + BT, T) - 1)*stride_qk + o_k + p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK) + b_gn = tl.load(p_gn, mask=m_k, other=0) + b_da += tl.load(p_dag, boundary_check=(0, 1)) * exp(b_ge) + b_dq += tl.load(p_dqg, boundary_check=(0, 1)) * exp(b_gi) * scale + tmp = exp(b_gn[None, :] - b_gi) + b_dk += tl.load(p_dkg, boundary_check=(0, 1)).to(tl.float32) * tmp + b_db += tl.load(p_dbg, boundary_check=(0, 1)).to(tl.float32) * tmp + tl.store(p_dq, (b_dq).to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b).to(tl.float32) + b_dgk_offset = b_da * b_a + tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_bwd_kernel_intra_tensorcore( + q, + k, + a, + b, + gi, + ge, + dAqk, + dAqb, + dAak, + dAab, + dq, + dk, + da, + db, + dqg, + dkg, + dag, + dbg, + dgk, + dgk_offset, + cu_seqlens, + chunk_indices, + scale: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + GATHER_SUPPORTED: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T_len = eos - bos + else: + bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + T_len = T + + if i_t * BT >= T_len: + return + + offset_base_k = (bos * H + i_h) * K + offset_base_attn = (bos * H + i_h) * BT + + valid_len = min(T_len - i_t * BT, BT) + mid_idx = valid_len // 2 + m_k = tl.arange(0, BK) + i_k * BK < K + p_offset = gi + offset_base_k + (i_t * BT + mid_idx) * K + tl.arange(0, BK) + i_k * BK + b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) + + # Q, K, A, B, Gates: [BT, BK] + p_q = tl.make_block_ptr(q + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_a = tl.make_block_ptr(a + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_b = tl.make_block_ptr(b + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_gi = tl.make_block_ptr(gi + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_ge = tl.make_block_ptr(ge + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) + b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + + b_gi_shifted = b_gi_val - b_offset[None, :] + b_ge_shifted = b_ge_val - b_offset[None, :] + exp_gi_shifted = tl.exp(b_gi_shifted) + inv_exp_gi_shifted = tl.exp(-b_gi_shifted) + exp_ge_shifted = tl.exp(b_ge_shifted) + + q_ops = (b_q * exp_gi_shifted).to(tl.float32) + k_ops = (b_k * inv_exp_gi_shifted).to(tl.float32) + b_ops = (b_b * inv_exp_gi_shifted).to(tl.float32) + a_ops = (b_a * exp_ge_shifted).to(tl.float32) + + p_dAqk = tl.make_block_ptr(dAqk + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) + p_dAqb = tl.make_block_ptr(dAqb + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) + p_dAak = tl.make_block_ptr(dAak + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) + p_dAab = tl.make_block_ptr(dAab + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) + + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) + b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) + b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) + + offs_n = tl.arange(0, BT) + offs_m = tl.arange(0, BT) + mask_inclusive = offs_m[:, None] >= offs_n[None, :] + mask_strict = offs_m[:, None] > offs_n[None, :] + + b_dAqk = tl.where(mask_inclusive, b_dAqk, 0.0).to(b_dAqk.dtype) + b_dAqb = tl.where(mask_inclusive, b_dAqb, 0.0).to(b_dAqb.dtype) + b_dAak = tl.where(mask_strict, b_dAak, 0.0).to(tl.float32) + b_dAab = tl.where(mask_strict, b_dAab, 0.0).to(tl.float32) + + # Intra-chunk gradients calculation + k_ops_h = k_ops.to(b_dAqk.dtype) + b_ops_h = b_ops.to(b_dAqb.dtype) + b_dq_intra = tl.dot(b_dAqk, k_ops_h) + tl.dot(b_dAqb, b_ops_h) + b_da_intra = tl.dot(b_dAak, k_ops) + tl.dot(b_dAab, b_ops) + + b_dAqk_T = tl.trans(b_dAqk) + b_dAqb_T = tl.trans(b_dAqb) + b_dAak_T = tl.trans(b_dAak) + b_dAab_T = tl.trans(b_dAab) + q_ops_h = q_ops.to(b_dAqk.dtype) + b_dk_intra = tl.dot(b_dAqk_T, q_ops_h) + tl.dot(b_dAak_T, a_ops) + b_db_intra = tl.dot(b_dAqb_T, q_ops_h) + tl.dot(b_dAab_T, a_ops) + + # Inter-chunk gradients loading + p_dqg = tl.make_block_ptr(dqg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dkg = tl.make_block_ptr(dkg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dag = tl.make_block_ptr(dag + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dbg = tl.make_block_ptr(dbg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + + b_dqg = tl.load(p_dqg, boundary_check=(0, 1)) + b_dag = tl.load(p_dag, boundary_check=(0, 1)) + b_dkg = tl.load(p_dkg, boundary_check=(0, 1)) + b_dbg = tl.load(p_dbg, boundary_check=(0, 1)) + + last_idx = min((i_t+1) * BT, T_len) - 1 + p_g_last = gi + offset_base_k + last_idx * H * K + tl.arange(0, BK) + i_k * BK + b_g_last = tl.load(p_g_last, mask=m_k, other=0.0) + + b_dq = b_dq_intra * exp_gi_shifted + b_dqg * tl.exp(b_gi_val) * scale + b_da = b_da_intra * exp_ge_shifted + b_dag * tl.exp(b_ge_val) + term_inter_stabilized = tl.exp(b_g_last[None, :] - b_gi_val) + b_dk = b_dk_intra * inv_exp_gi_shifted + b_dkg * term_inter_stabilized + b_db = b_db_intra * inv_exp_gi_shifted + b_dbg * term_inter_stabilized + + b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b) + b_dgk_offset = b_da * b_a + + p_dq = tl.make_block_ptr(dq + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_da = tl.make_block_ptr(da + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_db = tl.make_block_ptr(db + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dgk = tl.make_block_ptr(dgk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dgk_offset = tl.make_block_ptr(dgk_offset + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + for BK in [32, 64] + ], + key=['BK', 'BT', 'K'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_bwd_dgk_kernel( + dgk, + dgk_offset, + dgk_last, + dgk_output, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int32) + bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + + stride_qk = H * K + dgk += (bos * H + i_h) * K + dgk_offset += (bos * H + i_h) * K + dgk_last += (i_tg * H + i_h) * K + dgk_output += (bos * H + i_h) * K + p_dgk_last = dgk_last + tl.arange(0, BK) + i_k * BK + m_k = tl.arange(0, BK) + i_k * BK < K + b_dgk_last = tl.load(p_dgk_last, mask=m_k, other=0) + p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) + b_dgk_offset = tl.load(p_dgk_offset, boundary_check=(0, 1)) + # m_inv_cumsum = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]).to(tl.float32) + # b_dgk_cumsum = tl.dot(m_inv_cumsum, b_dgk, allow_tf32=False) + b_dgk_cumsum = tl.cumsum(b_dgk, 0, reverse=True) + b_dgk_cumsum += b_dgk_last[None, :] + b_dgk_cumsum -= b_dgk_offset + p_dgk_output = tl.make_block_ptr(dgk_output, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dgk_output, b_dgk_cumsum.to(p_dgk_output.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_dplr_bwd_dqk_intra( + q: torch.Tensor, + k: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + dAqk: torch.Tensor, + dAqb: torch.Tensor, + dAak: torch.Tensor, + dAab: torch.Tensor, + dqg: torch.Tensor, + dkg: torch.Tensor, + dag: torch.Tensor, + dbg: torch.Tensor, + dgk_last: torch.Tensor, + scale: float = 1.0, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = q.shape + BT = chunk_size + BK = min(64, triton.next_power_of_2(K)) if check_shared_mem() else min(32, triton.next_power_of_2(K)) + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NK = triton.cdiv(K, BK) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + da = torch.empty_like(a) + db = torch.empty_like(b) + dgk = torch.empty_like(gi, dtype=torch.float) + dgk_offset = torch.empty_like(gi, dtype=torch.float) + + grid = (NK, NT, B * H) + if safe_gate: + chunk_dplr_bwd_kernel_intra_func = chunk_dplr_bwd_kernel_intra_tensorcore + else: + chunk_dplr_bwd_kernel_intra_func = chunk_dplr_bwd_kernel_intra + chunk_dplr_bwd_kernel_intra_func[grid]( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + dAqk=dAqk, + dAqb=dAqb, + dAak=dAak, + dAab=dAab, + dq=dq, + dk=dk, + dgk=dgk, + dgk_offset=dgk_offset, + dqg=dqg, + dkg=dkg, + dag=dag, + dbg=dbg, + da=da, + db=db, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BT, + BK=BK, + GATHER_SUPPORTED=IS_GATHER_SUPPORTED, + ) + + dgk_output = torch.empty_like(dgk) + + def grid(meta): return (NT, triton.cdiv(K, meta['BK']), B * H) + chunk_dplr_bwd_dgk_kernel[grid]( + dgk=dgk, + dgk_offset=dgk_offset, + dgk_last=dgk_last, + dgk_output=dgk_output, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + ) + return dq, dk, da, db, dgk_output diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..026f3357eaa641006cefcad653fea0357b6a109a --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py @@ -0,0 +1,368 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp, gather +from fla.utils import IS_AMD, IS_GATHER_SUPPORTED, USE_CUDA_GRAPH, autotune_cache_kwargs + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BK', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_fwd_A_kernel_intra_sub_intra( + q, + k, + a, + b, + gi, + ge, + qg, + kg, + ag, + bg, + Aqk, + Aqb, + Aab, + Aak, + cu_seqlens, + chunk_indices, + scale: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + GATHER_SUPPORTED: tl.constexpr, +): + i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + tl.arange(0, BC)) < T + last_idx = min((i_t+1) * BT, T) - 1 + o_A = (bos + i_t * BT + tl.arange(0, BC)) * H*BT + i_h * BT + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_gi = tl.make_block_ptr(gi + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_ge = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_g_last = gi + (bos * H + i_h) * K + last_idx * H * K + tl.arange(0, BK) + b_g_last = tl.load(p_g_last, mask=m_k, other=0) + p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_bg = tl.make_block_ptr(bg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = b_q * scale + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_gi = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) + b_ge = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + + # deal with decay term. + g_exp = exp(b_gi) + g_exp_inv = exp(-b_gi + b_g_last[None, :]) + b_qg = b_q * g_exp + b_kg = b_k * g_exp_inv + b_bg = b_b * g_exp_inv + b_ag = b_a * exp(b_ge) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_bg, b_bg.to(p_bg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_ag, b_ag.to(p_ag.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + # tl.debug_barrier() + + b_q = b_q.to(b_k.dtype) + # inner attn + for j in range(0, min(BC, T - i_t * BT)): + # a trick to index the j-th row of b_k, b_g, b_b + if GATHER_SUPPORTED: + row_idx = tl.full([1, BK], j, dtype=tl.int16) + # [1, BK] + b_k_j = gather(b_k, row_idx, axis=0) + b_gk_j = gather(b_gi, row_idx, axis=0) + b_b_j = gather(b_b, row_idx, axis=0) + else: + mask = tl.arange(0, BC) == j + b_k_j = tl.sum(tl.where(mask[:, None], b_k, 0), 0)[None, :] + b_gk_j = tl.sum(tl.where(mask[:, None], b_gi, 0), 0)[None, :] + b_b_j = tl.sum(tl.where(mask[:, None], b_b, 0), 0)[None, :] + tmp = exp(b_gi - b_gk_j) + b_A_qk = tl.sum(b_q * b_k_j * tmp, 1) + m_i = (o_i >= j).to(tl.float32) + b_A_qk = b_A_qk * m_i + b_A_qb = tl.sum(b_q * b_b_j * tmp, 1) + b_A_qb = b_A_qb * m_i + tmp2 = exp(b_ge - b_gk_j) + b_A_ak = tl.sum(b_a * b_k_j * tmp2, 1) + m_i2 = (o_i > j).to(tl.float32) + b_A_ak = b_A_ak * m_i2 + b_A_ab = tl.sum(b_a * b_b_j * tmp2, 1) + b_A_ab = b_A_ab * m_i2 + + tl.store(Aqk + o_A + j, b_A_qk.to(dtype=Aqk.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) + tl.store(Aqb + o_A + j, b_A_qb.to(dtype=Aqb.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) + tl.store(Aab + o_A + j, b_A_ab.to(dtype=Aqb.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) + tl.store(Aak + o_A + j, b_A_ak.to(dtype=Aqk.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4, 8] + for num_stages in [2, 3] + ], + key=['BK', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_fwd_A_kernel_intra_tensorcore( + q, + k, + a, + b, + gi, + ge, + qg, + kg, + ag, + bg, + Aqk, + Aqb, + Aab, + Aak, + cu_seqlens, + chunk_indices, + scale: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + GATHER_SUPPORTED: tl.constexpr, +): + i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T_len = eos - bos + else: + bos = i_b * T + T_len = T + + if i_t * BT >= T_len: + return + + # Compute base offset for all tensors + offset_base = (bos * H + i_h) * K + + # Load the current chunk of Q, K, A, B and their gates + p_q = tl.make_block_ptr(q + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_a = tl.make_block_ptr(a + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_b = tl.make_block_ptr(b + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_gi = tl.make_block_ptr(gi + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_ge = tl.make_block_ptr(ge + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) + b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + + # Calculate the index of the middle element of the valid part of the chunk + valid_len = min(T_len - i_t * BT, BT) + mid_idx = valid_len // 2 + + # Load the offset vector from Global Memory + # p_offset points to gi[i_t*BT + mid_idx, :] + m_k = tl.arange(0, BK) < K + p_offset = gi + offset_base + (i_t * BT + mid_idx) * K + tl.arange(0, BK) + b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) + + # Apply offset to gate values + # These operations broadcast [BK] to [BT, BK] + b_gi_val = b_gi_val - b_offset[None, :] + b_ge_val = b_ge_val - b_offset[None, :] + + # Apply decay factors (now numerically safe) + # q_term ~ exp(gi - offset) + # k_term ~ exp(-gi + offset) + exp_gi = tl.exp(b_gi_val) + inv_exp_gi = tl.exp(-b_gi_val) + exp_ge = tl.exp(b_ge_val) + + b_q = (b_q * scale).to(tl.float32) + + # Compute gated operands for matrix multiplication + # q, k in l2 norm, < 1, i + q_ops = (b_q * exp_gi).to(tl.float32) + k_ops = (b_k * inv_exp_gi).to(tl.float32) + b_ops = (b_b * inv_exp_gi).to(tl.float32) + a_ops = (b_a * exp_ge).to(tl.float32) + + # Load gate values at the last position for inter-chunk decay + last_idx = min((i_t+1) * BT, T_len) - 1 + p_g_last = gi + offset_base + last_idx * H * K + tl.arange(0, BK) + b_g_last = tl.load(p_g_last, mask=m_k, other=0.0) + + exp_offset = tl.exp(b_offset) + b_g_centered = b_g_last - b_offset + exp_g_centered = tl.exp(b_g_centered) + + # Create pointers for writing + p_qg = tl.make_block_ptr(qg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_kg = tl.make_block_ptr(kg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_ag = tl.make_block_ptr(ag + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_bg = tl.make_block_ptr(bg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + + # Store gated Q and A + tl.store(p_qg, (q_ops * exp_offset[None, :]).to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_ag, (a_ops * exp_offset[None, :]).to(p_ag.dtype.element_ty), boundary_check=(0, 1)) + + # Store gated K and B + b_kg_g = k_ops * exp_g_centered[None, :] + b_bg_g = b_ops * exp_g_centered[None, :] + tl.store(p_kg, b_kg_g.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_bg, b_bg_g.to(p_bg.dtype.element_ty), boundary_check=(0, 1)) + + # Transpose K and B for dot product + k_ops_t = tl.trans(k_ops) + b_ops_t = tl.trans(b_ops) + + # Compute intra-chunk attention using TensorCores + q_ops_h = q_ops.to(b_q.dtype) + b_A_qk = tl.dot(q_ops_h, k_ops_t.to(b_q.dtype)) + b_A_qb = tl.dot(q_ops_h, b_ops_t.to(b_q.dtype)) + b_A_ak = tl.dot(a_ops, k_ops_t) + b_A_ab = tl.dot(a_ops, b_ops_t) + + # Create causal masks + offs_n = tl.arange(0, BT) + offs_m = tl.arange(0, BT) + mask_inclusive = offs_m[:, None] >= offs_n[None, :] + mask_strict = offs_m[:, None] > offs_n[None, :] + + # Apply causal masking + b_A_qk = tl.where(mask_inclusive, b_A_qk, 0.0) + b_A_qb = tl.where(mask_inclusive, b_A_qb, 0.0) + b_A_ak = tl.where(mask_strict, b_A_ak, 0.0) + b_A_ab = tl.where(mask_strict, b_A_ab, 0.0) + + # Store the intra-chunk attention matrices + offset_out_base = (bos * H + i_h) * BT + + p_Aqk = tl.make_block_ptr(Aqk + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_Aqb = tl.make_block_ptr(Aqb + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_Aak = tl.make_block_ptr(Aak + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_Aab = tl.make_block_ptr(Aab + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + tl.store(p_Aqk, b_A_qk.to(p_Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqb, b_A_qb.to(p_Aqb.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aak, b_A_ak.to(p_Aak.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aab, b_A_ab.to(p_Aab.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_dplr_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + scale: float, + chunk_size: int, + safe_gate: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = k.shape + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + Aqk = q.new_empty(B, T, H, BT, dtype=q.dtype) + Aqb = q.new_empty(B, T, H, BT, dtype=q.dtype) + # involving matrix inverse and it'd be better to use float here. + Aab = q.new_empty(B, T, H, BT, dtype=torch.float) + Aak = q.new_empty(B, T, H, BT, dtype=torch.float) + + grid = (NT, B, H) + BK = max(triton.next_power_of_2(K), 16) + qg = torch.empty_like(q) + kg = torch.empty_like(k, dtype=q.dtype) + ag = torch.empty_like(a, dtype=q.dtype) + bg = torch.empty_like(b, dtype=q.dtype) + if safe_gate: + chunk_dplr_fwd_A_kernel_intra_func = chunk_dplr_fwd_A_kernel_intra_tensorcore + else: + chunk_dplr_fwd_A_kernel_intra_func = chunk_dplr_fwd_A_kernel_intra_sub_intra + chunk_dplr_fwd_A_kernel_intra_func[grid]( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + Aqk=Aqk, + Aqb=Aqb, + Aab=Aab, + Aak=Aak, + qg=qg, + kg=kg, + ag=ag, + bg=bg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BT, + BK=BK, + GATHER_SUPPORTED=IS_GATHER_SUPPORTED, + ) + return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..4451b95d907ae7b9905a89efbcc7db4162c8a04b --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py @@ -0,0 +1,175 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import IS_AMD, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +@triton.heuristics({ + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BT', 'BK', 'BV', "V"], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_bwd_kernel_dhu( + qg, + bg, + w, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1)) + + mask_k = tl.arange(0, BK) < K + for i_t in range(NT - 1, -1, -1): + p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + b_dh_tmp = tl.zeros([BK, BV], dtype=tl.float32) + for i_c in range(tl.cdiv(BT, BC) - 1, -1, -1): + p_qg = tl.make_block_ptr(qg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) + p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + # [BK, BT] + b_qg = tl.load(p_qg, boundary_check=(0, 1)) + # [BT, BK] + b_bg = tl.load(p_bg, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + # [BT, V] + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dv2 = b_dv + tl.dot(b_bg, b_dh.to(b_bg.dtype)) + tl.store(p_dv2, b_dv2.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + # [BK, BV] + b_dh_tmp += tl.dot(b_qg, b_do.to(b_qg.dtype)) + b_dh_tmp += tl.dot(b_w, b_dv2.to(b_qg.dtype)) + last_idx = min((i_t + 1) * BT, T) - 1 + bg_last = tl.load(gk + ((bos + last_idx) * H + i_h) * K + tl.arange(0, BK), mask=mask_k) + b_dh *= exp(bg_last)[:, None] + b_dh += b_dh_tmp + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_dplr_bwd_dhu( + qg: torch.Tensor, + bg: torch.Tensor, + w: torch.Tensor, + gk: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor | None, + do: torch.Tensor, + dv: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *qg.shape, do.shape[-1] + BT = chunk_size + BK = max(triton.next_power_of_2(K), 16) + assert BK <= 256, "current kernel does not support head dimension being larger than 256." + # H100 + if check_shared_mem('hopper', qg.device.index): + BV = 64 + BC = 64 if K <= 128 else 32 + elif check_shared_mem('ampere', qg.device.index): # A100 + BV = 32 + BC = 32 + else: # Etc: 4090 + BV = 16 + BC = 16 + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + BC = min(BT, BC) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + + dh = qg.new_empty(B, NT, H, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.zeros_like(dv) + + grid = (NK, NV, N * H) + chunk_dplr_bwd_kernel_dhu[grid]( + qg=qg, + bg=bg, + w=w, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BC=BC, + BK=BK, + BV=BV, + ) + return dh, dh0, dv2 diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..f8c0611c01e1aee616390a5c617f941c99fbe945 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py @@ -0,0 +1,175 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import IS_AMD, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BT', 'BK', 'BV'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_fwd_kernel_h( + kg, + v, + w, + bg, + u, + v_new, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + o_k = i_k * BK + tl.arange(0, BK) + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + + b_hc = tl.zeros([BK, BV], dtype=tl.float32) + # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden + for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): + p_kg = tl.make_block_ptr(kg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) + # [BK, BC] + b_kg = tl.load(p_kg, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_bg = tl.load(p_bg, boundary_check=(0, 1)) + b_v2 = tl.dot(b_w, b_h.to(b_w.dtype)) + tl.load(p_u, boundary_check=(0, 1)) + b_hc += tl.dot(b_kg, b_v) + b_hc += tl.dot(b_bg.to(b_hc.dtype), b_v2) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k, mask=o_k < K).to(tl.float32) + b_h *= exp(b_g_last[:, None]) + b_h += b_hc + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + + +def chunk_dplr_fwd_h( + kg: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + bg: torch.Tensor, + gk: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *kg.shape, u.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + BK = max(triton.next_power_of_2(K), 16) + assert BK <= 256, "current kernel does not support head dimension larger than 256." + # H100 can have larger block size + + if check_shared_mem('hopper', kg.device.index): + BV = 64 + BC = 64 if K <= 128 else 32 + elif check_shared_mem('ampere', kg.device.index): # A100 + BV = 32 + BC = 32 + else: + BV = 16 + BC = 16 + + BC = min(BT, BC) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + + h = kg.new_empty(B, NT, H, K, V) + final_state = kg.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + v_new = torch.empty_like(u) + grid = (NK, NV, N * H) + chunk_dplr_fwd_kernel_h[grid]( + kg=kg, + v=v, + w=w, + bg=bg, + u=u, + v_new=v_new, + h=h, + gk=gk, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BC=BC, + BK=BK, + BV=BV, + ) + return h, v_new, final_state diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..0c02ff0b4c45d8cc558fc6bf00f8155c785ecdf6 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py @@ -0,0 +1,436 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import IS_AMD, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + +BK_LIST = [32, 64, 128] if check_shared_mem() else [16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BV', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_bwd_kernel_dAu( + v, + do, + v_new, + A_qb, + dA_qk, + dA_qb, + dv_new, + cu_seqlens, + chunk_indices, + scale: tl.constexpr, + T, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + + b_dA_qk = tl.zeros([BT, BT], dtype=tl.float32) + b_dA_qb = tl.zeros([BT, BT], dtype=tl.float32) + + p_A_qb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + b_A_qb = tl.load(p_A_qb, boundary_check=(0, 1)) + # causal mask + b_A_qb = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_A_qb, 0.).to(b_A_qb.dtype) + + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_v_new = tl.make_block_ptr(v_new + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_dv_new = tl.make_block_ptr(dv_new + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_dA_qk += tl.dot(b_do, b_v) + b_dA_qb += tl.dot(b_do, b_v_new) + b_dv_new = tl.dot(tl.trans(b_A_qb), b_do) + # for recurrent + tl.store(p_dv_new, b_dv_new.to(p_dv_new.dtype.element_ty), boundary_check=(0, 1)) + + p_dA_qk = tl.make_block_ptr(dA_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dA_qb = tl.make_block_ptr(dA_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + b_dA_qk = tl.where(m_s, b_dA_qk * scale, 0.) + tl.store(p_dA_qk, b_dA_qk.to(p_dA_qk.dtype.element_ty), boundary_check=(0, 1)) + b_dA_qb = tl.where(m_s, b_dA_qb * scale, 0.) + tl.store(p_dA_qb, b_dA_qb.to(p_dA_qb.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BT', 'BK', 'BV'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def chunk_dplr_bwd_o_kernel( + v, + v_new, + h, + do, + dh, + dk, + db, + w, + dq, + dv, + dw, + gk, + dgk_last, + k, + b, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h += (i_tg * H + i_h) * K * V + dh += (i_tg * H + i_h) * K * V + dk += (bos * H + i_h) * K + k += (bos * H + i_h) * K + db += (bos * H + i_h) * K + b += (bos * H + i_h) * K + dw += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dq += (bos * H + i_h) * K + w += (bos * H + i_h) * K + + dgk_last += (i_tg * H + i_h) * K + gk += (bos * H + i_h) * K + + stride_qk = H*K + stride_vo = H*V + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_db = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk_last = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dgk_last += tl.sum((b_h * b_dh).to(tl.float32), axis=0) + + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + b_db += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) + p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) + + m_k = (i_k*BK+tl.arange(0, BK)) < K + last_idx = min(i_t * BT + BT, T) - 1 + b_gk_last = tl.load(gk + last_idx * stride_qk + i_k*BK + tl.arange(0, BK), mask=m_k, other=float('-inf')) + b_dgk_last *= exp(b_gk_last) + p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_dgk_last += tl.sum(b_k * b_dk, axis=0) + b_dgk_last += tl.sum(b_b * b_db, axis=0) + tl.store(dgk_last + tl.arange(0, BK) + i_k * BK, b_dgk_last, mask=m_k) + + p_dw = tl.make_block_ptr(dw, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + for BK in BK_LIST + for BV in BK_LIST + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def chunk_dplr_bwd_kernel_dv( + A_qk, + kg, + do, + dv, + dh, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + # offset calculation + A_qk += (bos * H + i_h) * BT + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + kg += (bos * H + i_h) * K + dh += (i_tg * H + i_h) * K*V + + stride_qk = H*K + stride_vo = H*V + stride_A = H*BT + + for i_k in range(tl.cdiv(K, BK)): + p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_kg = tl.make_block_ptr(kg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_kg = tl.load(p_kg, boundary_check=(0, 1)) + b_dv += tl.dot(b_kg, b_dh.to(b_kg.dtype)) + + p_Aqk = tl.make_block_ptr(A_qk, (BT, T), (1, stride_A), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], tl.load(p_Aqk, boundary_check=(0, 1)), 0) + p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv += tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_dplr_bwd_dv( + A_qk: torch.Tensor, + kg: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K, V = *kg.shape, do.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dv = torch.empty_like(do) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_dplr_bwd_kernel_dv[grid]( + A_qk=A_qk, + kg=kg, + do=do, + dv=dv, + dh=dh, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dv + + +def chunk_dplr_bwd_o( + k: torch.Tensor, + b: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + gk: torch.Tensor, + do: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + w: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + scale: float = 1.0, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + B, T, H, K, V = *w.shape, v.shape[-1] + + BT = chunk_size + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BK = min(max(triton.next_power_of_2(K), 16), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) + BV = min(max(triton.next_power_of_2(V), 16), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) + NK = triton.cdiv(K, BK) + dq = torch.empty_like(k) + dk = torch.empty_like(k) + dw = torch.empty_like(w) + db = torch.empty_like(b) + grid = (NK, NT, B * H) + + dgk_last = torch.empty(B, NT, H, K, dtype=torch.float, device=w.device) + + chunk_dplr_bwd_o_kernel[grid]( + k=k, + b=b, + v=v, + v_new=v_new, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + db=db, + dgk_last=dgk_last, + w=w, + dv=dv, + dw=dw, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dq, dk, dw, db, dgk_last + + +def chunk_dplr_bwd_dAu( + v: torch.Tensor, + v_new: torch.Tensor, + do: torch.Tensor, + A_qb: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, V = v.shape + BT = chunk_size + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + if check_shared_mem('ampere'): # A100 + BV = min(triton.next_power_of_2(V), 128) + elif check_shared_mem('ada'): # 4090 + BV = min(max(triton.next_power_of_2(V), 16), 64) + else: + BV = min(triton.next_power_of_2(V), 32) + + grid = (NT, B * H) + dA_qk = torch.empty(B, T, H, BT, dtype=torch.float, device=v.device) + dA_qb = torch.empty(B, T, H, BT, dtype=torch.float, device=v.device) + dv_new = torch.empty_like(v_new) + chunk_dplr_bwd_kernel_dAu[grid]( + v=v, + do=do, + v_new=v_new, + A_qb=A_qb, + dA_qk=dA_qk, + dA_qb=dA_qb, + dv_new=dv_new, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + V=V, + BT=BT, + BV=BV, + ) + return dv_new, dA_qk, dA_qb diff --git a/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py b/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a05730743deb9dd84c46d2b6f9783878fd88b5 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py @@ -0,0 +1,125 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_AMD, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + +BK_LIST = [32, 64, 128] if check_shared_mem() else [16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BK_LIST + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3, 4] + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_dplr_fwd_kernel_o( + qg, + v, + v_new, + A_qk, + A_qb, + h, + o, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_qg = tl.load(p_qg, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_o += tl.dot(b_qg, b_h) + + p_Aqk = tl.make_block_ptr(A_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_Aqb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + b_Aqk = tl.load(p_Aqk, boundary_check=(0, 1)) + b_Aqb = tl.load(p_Aqb, boundary_check=(0, 1)) + b_Aqk = tl.where(m_s, b_Aqk, 0) + b_Aqb = tl.where(m_s, b_Aqb, 0) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_o = b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_v_new.dtype), b_v_new) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_dplr_fwd_o( + qg: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + A_qk: torch.Tensor, + A_qb: torch.Tensor, + h: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K, V = *qg.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = torch.empty_like(v) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_dplr_fwd_kernel_o[grid]( + qg=qg, + v=v, + v_new=v_new, + A_qk=A_qk, + A_qb=A_qb, + h=h, + o=o, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o diff --git a/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..9248ee1c8bf021fc728faac520a65c9cf3804956 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/fused_recurrent.py @@ -0,0 +1,266 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import USE_CUDA_GRAPH, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BV in [16, 32, 64] + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=['BK'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_dplr_delta_rule_fwd_kernel( + q, + k, + v, + a, + b, + gk, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_a = a + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_b = b + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_gk = gk + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_o = o + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) + b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + + b_h = exp(b_gk)[:, None] * b_h + b_b[:, None] * tl.sum(b_a[:, None] * b_h, 0)[None, :] + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + p_q += (-1 if REVERSE else 1) * H*K + p_k += (-1 if REVERSE else 1) * H*K + p_a += (-1 if REVERSE else 1) * H*K + p_b += (-1 if REVERSE else 1) * H*K + p_gk += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_o += (-1 if REVERSE else 1) * H*V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_dplr_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float | None = 1.0, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + + h0 = initial_state + ht = q.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + o = torch.empty_like(v) + + def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) + fused_recurrent_dplr_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + a=a, + b=b, + gk=gk, + o=o, + h0=h0, + ht=ht, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + REVERSE=reverse, + ) + return o, ht + + +class FusedRecurrentDPLRDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, ht = fused_recurrent_dplr_delta_rule_fwd( + q=q, + k=k, + v=v, + a=a, + b=b, + gk=gk, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + return o, ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass for fused_recurrent_dplr_delta_rule is not implemented and will not be supported. " + "This kernel is only for inference. " + "For training, please use `chunk_dplr_delta_rule`.", + ) + + +def fused_recurrent_dplr_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + This function computes the recurrence S_t = S_t @ (Diag(g_t) + a_t b_t^T) + v_t k_t^T in a recurrent manner. + + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, T, H, K]`. + gk (torch.Tensor): + gk of shape `[B, T, H, K]`. decay term in log space! + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (Optional[torch.Tensor]): + Cumulative sequence lengths of shape `[N + 1]` used for variable-length training, + consistent with the FlashAttention API. + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = q.shape[-1] ** -0.5 + o, final_state = FusedRecurrentDPLRDeltaRuleFunction.apply( + q, + k, + v, + a, + b, + gk, + scale, + initial_state, + output_final_state, + reverse, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/generalized_delta_rule/dplr/naive.py b/fla/ops/generalized_delta_rule/dplr/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..1abfdf8ea1267f71319b9b6370e841f5f081ff39 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/naive.py @@ -0,0 +1,95 @@ + +import torch +from einops import rearrange + +# S_t = S_t @ (I + alpha_t beta_t^T) + v_t k_t^T +# q, k, alpha, beta [B, H, L, D_K] +# v [B, H, L, D_V] + + +def dplr_recurrence(q, k, v, alpha, beta, gk, initial_state=None, output_final_state=True): + orig_dtype = q.dtype + b, h, l, d_k = q.shape + q, k, v, beta, gk = map(lambda x: x.float(), [q, k, v, beta, gk]) + d_v = v.shape[-1] + o = torch.zeros_like(v) + S = torch.zeros(b, h, d_k, d_v).to(v) + q = q * (d_k ** -0.5) + + if initial_state is not None: + S += initial_state + + for i in range(l): + _k = k[:, :, i] + _q = q[:, :, i] + _v = v[:, :, i] + _alpha = alpha[:, :, i].clone() + _beta = beta[:, :, i].clone() + _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _alpha[..., None]).sum(-2, keepdim=True) * _beta[..., None] + S = S.clone() * gk[:, :, i].exp()[..., None] + _kv + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) + S = None if output_final_state is False else S + return o.to(orig_dtype), S + + +def dplr_chunkwise(q, k, v, alpha, beta, gk, initial_state=None, output_final_state=True, chunk_size=32): + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * (d_k ** -0.5) + v = v + assert l % chunk_size == 0 + + S = k.new_zeros(b, h, d_k, d_v).to(q) + if initial_state is not None: + S += initial_state + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, alpha, beta, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', + c=chunk_size).float(), [q, k, v, alpha, beta, gk]) + + gk_cumsum = gk.cumsum(-2) + + # v2 = (alpha @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v + A_ab = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) + A_qk = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) + A_ak = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) + A_qb = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) + + for i in range(chunk_size): + alpha_i = alpha[:, :, :, i, None] + q_i = q[:, :, :, i, None] + gk_i = gk_cumsum[:, :, :, i, None] + mask = (torch.arange(chunk_size) <= i).to(q.device) + attn_i = (gk_i - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() + A_qk[:, :, :, i, :] = (q_i * k * attn_i).sum(-1).clone() + A_qb[:, :, :, i, :] = (q_i * beta * attn_i).sum(-1).clone() + mask = (torch.arange(chunk_size) < i).to(q.device) + # shift by one. + attn_i = (gk_i - gk[:, :, :, i, None] - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() + A_ab[:, :, :, i, :] = (alpha_i * beta * attn_i).sum(-1).clone() + A_ak[:, :, :, i, :] = (alpha_i * k * attn_i).sum(-1).clone() + + A_ab = A_ab + for i in range(1, chunk_size): + A_ab[..., i, :i] = A_ab[..., i, :i].clone() + (A_ab[..., i, :, None].clone() * A_ab[..., :, :i].clone()).sum(-2) + + A_ab = A_ab + torch.eye(chunk_size, dtype=torch.float, device=q.device) + u = A_ab @ (A_ak @ v) + w = A_ab @ ((gk_cumsum-gk).exp() * alpha) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i, u_i, w_i, beta_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], beta[:, :, i] + v2_i = u_i + w_i @ S + + o_1 = A_qk[:, :, i] @ v_i + o_2 = A_qb[:, :, i] @ v2_i + o_3 = (q_i * gk_cumsum[:, :, i].exp()) @ S + o[:, :, i] = o_1 + o_2 + o_3 + decay = (gk_cumsum[:, :, i, -1, None] - gk_cumsum[:, :, i]).exp() + S = S*gk_cumsum[:, :, i, -1, :, None].exp() + (k_i * decay).transpose(-1, -2) @ v_i + \ + (beta_i * decay).transpose(-1, -2) @ v2_i + S = None if output_final_state is False else S + return rearrange(o, 'b h n c d -> b h (n c) d'), S diff --git a/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py b/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..5ba4dcca339d26fc4d23a30425ac082d8b12b282 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py @@ -0,0 +1,164 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_INTEL_ALCHEMIST, USE_CUDA_GRAPH, autotune_cache_kwargs, check_shared_mem + +# https://github.com/intel/intel-xpu-backend-for-triton/issues/3449 +triton_config = {'grf_mode': 'large'} if IS_INTEL_ALCHEMIST else {} + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config(triton_config, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=['BT', 'BK', 'BV'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + A_ab_inv, + A_ak, + ag, + v, + dw, + du, + dv, + dv0, + dag, + dAak, + dAab, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_Aak_t = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_Aab_inv_t = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_dAak = tl.make_block_ptr(dAak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dAab = tl.make_block_ptr(dAab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + b_A_ab_inv_t = tl.load(p_Aab_inv_t, boundary_check=(0, 1)) + b_A_ak_t = tl.load(p_Aak_t, boundary_check=(0, 1)) + b_A_ak_t = tl.where(tl.arange(0, BT)[:, None] < tl.arange(0, BT)[None, :], b_A_ak_t, 0) + b_A_ab_inv_t = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A_ab_inv_t, 0) + b_A_tmp_t = tl.dot(b_A_ak_t, b_A_ab_inv_t).to(v.dtype.element_ty) + b_dA_tmp = tl.zeros([BT, BT], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv0 = tl.make_block_ptr(dv0 + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA_tmp += tl.dot(b_du.to(b_v.dtype), tl.trans(b_v)) + b_dv0 = tl.load(p_dv0, boundary_check=(0, 1)) + b_dv = b_dv0 + tl.dot(b_A_tmp_t, b_du) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + m_i = tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :] + b_dA_tmp = tl.where(m_i, b_dA_tmp, 0) + b_dA_ak = tl.dot(b_A_ab_inv_t, b_dA_tmp) + b_dA_ak = tl.where(m_i, b_dA_ak, 0) + tl.store(p_dAak, b_dA_ak, boundary_check=(0, 1)) + b_dA_ab_inv = tl.dot(b_dA_tmp, b_A_ak_t) + + for i_k in range(tl.cdiv(K, BK)): + p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dag = tl.make_block_ptr(dag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_ag = tl.load(p_ag, boundary_check=(0, 1)) + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dA_ab_inv += tl.dot(b_dw, tl.trans(b_ag)) + b_dag = tl.dot(b_A_ab_inv_t.to(b_dw.dtype), b_dw) + tl.store(p_dag, b_dag.to(p_dag.dtype.element_ty), boundary_check=(0, 1)) + + # if we know dL/dA^(-1), for dL/dA, we can use the following formula: + # dL/dA = -(A^(-1))^T @ (dL/dA^(-1)) @ (A^(-1))^T + # in the fwd pass we use fwd substitution to calculate (I-lower(A_ab))^-1. + # denote A = I - lower(A_ab), B = A^-1 + # in the backward pass. + # dL/dA = -(B)^T @ (dL/dB) @ B^T + # dL/dA_ab = lower(B^T @ dL/dB @ B^T) + b_dA_ab_inv = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dA_ab_inv, 0) + b_dA_ab_inv = tl.dot(b_A_ab_inv_t, b_dA_ab_inv) + b_dA_ab_inv = tl.dot(b_dA_ab_inv, b_A_ab_inv_t) + b_dA_ab_inv = tl.where(m_i, b_dA_ab_inv, 0) + tl.store(p_dAab, b_dA_ab_inv, boundary_check=(0, 1)) + + +def chunk_dplr_bwd_wy( + A_ab_inv: torch.Tensor, + A_ak: torch.Tensor, + v: torch.Tensor, + ag: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + dv0: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + A_ab_inv, A_ak, v, ag, dw, du = map(lambda x: x.contiguous(), [A_ab_inv, A_ak, v, ag, dw, du]) + B, T, H, K, V = *dw.shape, du.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = min(max(triton.next_power_of_2(K), 16), 64) + BV = min(max(triton.next_power_of_2(V), 16), 64) if check_shared_mem() else min(max(triton.next_power_of_2(V), 16), 32) + + dA_ab = torch.empty_like(A_ab_inv, dtype=torch.float) + dA_ak = torch.empty_like(A_ak, dtype=torch.float) + dv = torch.empty_like(v) + dag = torch.empty_like(ag) + + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + A_ab_inv=A_ab_inv, + A_ak=A_ak, + ag=ag, + v=v, + dw=dw, + du=du, + dv=dv, + dv0=dv0, + dag=dag, + dAak=dA_ak, + dAab=dA_ab, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dA_ab, dA_ak, dv, dag diff --git a/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py b/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..05f61a0a79a7e804a6912cd559379c21a4c8773d --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py @@ -0,0 +1,289 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import gather +from fla.utils import IS_GATHER_SUPPORTED, USE_CUDA_GRAPH, autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_fwd_kernel_chunk32( + A_ab, + A_ab_inv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, # placeholder, do not delete + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_Aab = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_Aab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A_ab = tl.load(p_Aab, boundary_check=(0, 1)) + b_A_ab = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A_ab, 0) + for i in range(1, BT): + mask = tl.arange(0, BT) == i + b_a = tl.sum(tl.where(mask[:, None], b_A_ab, 0), 0) + b_a = b_a + tl.sum(b_a[:, None] * b_A_ab, 0) * (tl.arange(0, BT) < i) + b_A_ab = tl.where(mask[:, None], b_a, b_A_ab) + b_A_ab += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] + tl.store(p_Aab_inv, b_A_ab.to(p_Aab_inv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BC'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_fwd_kernel_chunk64( + A_ab, + A_ab_inv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + IS_VARLEN: tl.constexpr, + GATHER_SUPPORTED: tl.constexpr = IS_GATHER_SUPPORTED, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_A1 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) + p_A2 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) + p_A3 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) + p_A_inv1 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) + p_A_inv2 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) + p_A_inv3 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) + p_A_inv4 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) + + b_A = tl.load(p_A1, boundary_check=(0, 1)) + b_A2 = tl.load(p_A2, boundary_check=(0, 1)) + b_A3 = tl.load(p_A3, boundary_check=(0, 1)) + b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) + b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) + + for i in range(1, BC): + if GATHER_SUPPORTED: + row_idx = tl.full([1, BC], i, dtype=tl.int16) + # [1, BK] -> [BK] + b_a = tl.sum(gather(b_A, row_idx, axis=0), 0) + b_a2 = tl.sum(gather(b_A2, row_idx, axis=0), 0) + else: + mask = tl.arange(0, BC) == i + b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) + b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) + mask = tl.arange(0, BC) == i + # b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) + # b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BC) < i) + b_a2 = b_a2 + tl.sum(b_a2[:, None] * b_A2, 0) * (tl.arange(0, BC) < i) + b_A = tl.where(mask[:, None], b_a, b_A) + b_A2 = tl.where(mask[:, None], b_a2, b_A2) + + # blockwise computation of lower triangular matrix's inverse + # i.e., [A11, 0; A21, A22]^-1 = [A11^-1, 0; -A22^-1 A21 A11^-1, A22^-1] + b_A += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] + b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] + b_A3 = tl.dot(tl.dot(b_A2, b_A3), b_A) + # tl.debug_barrier() + tl.store(p_A_inv1, b_A.to(p_A_inv1.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_A_inv2, b_A2.to(p_A_inv2.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_A_inv3, b_A3.to(p_A_inv3.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + # causal mask + tl.store(p_A_inv4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A_inv4.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def wu_fwd_kernel( + w, + u, + ag, + v, + A_ab_inv, + A_ak, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_s = tl.arange(0, BT) + + p_A_ab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_A_ak = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + b_Aab_inv = tl.load(p_A_ab_inv, boundary_check=(0, 1)) + b_Aak = tl.load(p_A_ak, boundary_check=(0, 1)) + b_Aab_inv = tl.where(o_s[:, None] >= o_s[None, :], b_Aab_inv, 0) + b_Aak = tl.where(o_s[:, None] > o_s[None, :], b_Aak, 0) + # let's use tf32 here + b_Aak = tl.dot(b_Aab_inv, b_Aak) + # (SY 01/04) should be bf16 or tf32? To verify. + b_Aak = b_Aak.to(v.dtype.element_ty, fp_downcast_rounding="rtne") + b_Aab_inv = b_Aab_inv.to(ag.dtype.element_ty, fp_downcast_rounding="rtne") + + for i_k in range(tl.cdiv(K, BK)): + p_ag = tl.make_block_ptr(ag + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_ag = tl.load(p_ag, boundary_check=(0, 1)) + b_w = tl.dot(b_Aab_inv, b_ag) # both bf16 or fp16 + tl.store(p_w, b_w.to(p_w.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_u = tl.dot(b_Aak, b_v) # both bf16 or fp16 + tl.store(p_u, b_u.to(p_u.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + + +def wu_fwd( + ag: torch.Tensor, + v: torch.Tensor, + A_ak: torch.Tensor, + A_ab_inv: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *ag.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = min(max(triton.next_power_of_2(K), 16), 64) + BV = min(max(triton.next_power_of_2(V), 16), 64) + + w = torch.empty_like(ag) + u = torch.empty_like(v) + wu_fwd_kernel[(NT, B * H)]( + ag=ag, + v=v, + A_ak=A_ak, + A_ab_inv=A_ab_inv, + w=w, + u=u, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +def prepare_wy_repr_fwd( + ag: torch.Tensor, + v: torch.Tensor, + A_ak: torch.Tensor, + A_ab: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, _ = ag.shape + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BC = min(BT, 32) + fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 + A_ab_inv = torch.empty_like(A_ab) + fwd_fn[(NT, B * H)]( + A_ab=A_ab, + A_ab_inv=A_ab_inv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + BC=BC, + ) + w, u = wu_fwd( + ag=ag, + v=v, + A_ak=A_ak, + A_ab_inv=A_ab_inv, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + return w, u, A_ab_inv + + +fwd_prepare_wy_repr = prepare_wy_repr_fwd + +fwd_wu = wu_fwd diff --git a/fla/ops/generalized_delta_rule/iplr/__init__.py b/fla/ops/generalized_delta_rule/iplr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e1b0916fc99cb5783fe959ae5d40b89bc9326425 --- /dev/null +++ b/fla/ops/generalized_delta_rule/iplr/__init__.py @@ -0,0 +1,7 @@ +from .chunk import chunk_iplr_delta_rule +from .fused_recurrent import fused_recurrent_iplr_delta_rule + +__all__ = [ + 'chunk_iplr_delta_rule', + 'fused_recurrent_iplr_delta_rule', +] diff --git a/fla/ops/generalized_delta_rule/iplr/chunk.py b/fla/ops/generalized_delta_rule/iplr/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..bef406f2f9ca363de6c2bdf3e826c0c92a3281ad --- /dev/null +++ b/fla/ops/generalized_delta_rule/iplr/chunk.py @@ -0,0 +1,512 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.generalized_delta_rule.iplr.wy_fast import prepare_wy_repr_fwd +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.utils import ( + USE_CUDA_GRAPH, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_shared_mem, + input_guard, +) + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4] + ([] if check_shared_mem('hopper') else [8]) + ], + key=['BT', 'BK', 'BV'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_generalized_iplr_delta_rule_fwd_kernel_h( + k, + v, + d, + b, + u, + v_new, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + b_hc = tl.zeros([BK, BV], dtype=tl.float32) + # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden + for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_b = tl.make_block_ptr(b+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) + p_d = tl.make_block_ptr(d+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_d = tl.load(p_d, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_v2 = tl.dot(b_d, b_h.to(b_d.dtype)) + tl.load(p_u, boundary_check=(0, 1)) + b_hc += tl.dot(b_k, b_v) + b_hc += tl.dot(b_b, b_v2.to(b_k.dtype)) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + b_h += b_hc + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in [2, 4, 8] + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_generalized_iplr_delta_rule_fwd_kernel_o( + q, + k, + v, + u, + b, + h, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + b += (bos * H + i_h) * K + v += (bos * H + i_h) * V + u += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h) * K * V + stride_qk = H*K + stride_vo = H*V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_Aqk = tl.zeros([BT, BT], dtype=tl.float32) + b_Aqb = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_b = tl.make_block_ptr(b, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_Aqk += tl.dot(b_q, b_k) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_Aqb += tl.dot(b_q, b_b) + + o_i = tl.arange(0, BT) + m_A = o_i[:, None] >= o_i[None, :] + b_Aqk = tl.where(m_A, b_Aqk, 0) + b_Aqb = tl.where(m_A, b_Aqb, 0) + + p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_u = tl.load(p_u, boundary_check=(0, 1)) + b_o = (b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_u.dtype), b_u)) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_generalized_iplr_delta_rule_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + b: torch.Tensor, + h: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K, V = *q.shape, v.shape[-1] + if scale is None: + scale = k.shape[-1] ** -0.5 + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = torch.empty_like(v) + + def grid(meta): return ( + triton.cdiv(V, meta['BV']), + NT, + B * H, + ) + chunk_generalized_iplr_delta_rule_fwd_kernel_o[grid]( + q=q, + k=k, + v=v, + u=v_new, + b=b, + h=h, + o=o, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +def chunk_generalized_iplr_delta_rule_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + b: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, u.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + BK = max(triton.next_power_of_2(K), 16) + assert BK <= 256, "current kernel does not support head dimension larger than 256." + # H100 can have larger block size + + if check_shared_mem('hopper', k.device.index): + BV = 64 + BC = 64 if K <= 128 else 32 + elif check_shared_mem('ampere', k.device.index): # A100 + BV = 32 + BC = 32 + else: + BV = 16 + BC = 16 + + BC = min(BT, BC) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) + grid = (NK, NV, N * H) + + chunk_generalized_iplr_delta_rule_fwd_kernel_h[grid]( + k=k, + v=v, + d=w, + b=b, + u=u, + v_new=v_new, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BC=BC, + BK=BK, + BV=BV, + ) + return h, v_new, final_state + + +def chunk_generalized_iplr_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + w, u, _ = prepare_wy_repr_fwd( + a=a, + b=b, + k=k, + v=v, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + h, v_new, final_state = chunk_generalized_iplr_delta_rule_fwd_h( + k=k, + v=v, + b=b, + w=w, + u=u, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + o = chunk_generalized_iplr_delta_rule_fwd_o( + q=q, + k=k, + v=v, + v_new=v_new, + b=b, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return o, final_state + + +class ChunkGeneralizedIPLRDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + ): + chunk_size = min(64, max(triton.next_power_of_2(q.shape[1]), 16)) + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + o, final_state = chunk_generalized_iplr_delta_rule_fwd( + q=q, + k=k, + v=v, + a=a, + b=b, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + raise NotImplementedError( + "Backward pass for ChunkGeneralizedIPLRDeltaRuleFunction is not implemented yet. " + "Stay tuned!", + ) + + +@torch.compiler.disable +def chunk_iplr_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + a (torch.Tensor): + activations of shape `[B, T, H, K]`. + b (torch.Tensor): + betas of shape `[B, T, H, K]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cu_seqlens_cpu (torch.LongTensor): + CPU version of cumulative sequence lengths of shape `[N+1]`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + scale = k.shape[-1] ** -0.5 if scale is None else scale + o, final_state = ChunkGeneralizedIPLRDeltaRuleFunction.apply( + q, + k, + v, + a, + b, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..3a8596840f8ee7f06222ffa4857523e20220d152 --- /dev/null +++ b/fla/ops/generalized_delta_rule/iplr/fused_recurrent.py @@ -0,0 +1,452 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.utils import autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BV in [32, 64] + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=["BK"], + **autotune_cache_kwargs, +) +@triton.jit +def fused_recurrent_fwd_kernel( + q, # query [B, H, L, K] + k, # key [B, H, L, V] + v, # value [B, H, L, V]. + a, # a [B, H, L, K] + b, # b [B, H, L, K] + o, # output [B, H, L, V] + ha, # tmp variable [B, H, L, V] for storing intermediate results of (h * a[None, :]).sum(0) + h0, # initial hidden state [B, H, K, V] + ht, # final hidden state [B, H, K, V] + cu_seqlens, # varlen cu_seqlens + scale, # K ** -0.5 + H, # n_heads + T, # seq_len + K: tl.constexpr, # K + V: tl.constexpr, # V + BK: tl.constexpr, # BLOCK SIZE along the K dimension + BV: tl.constexpr, # BLOCK SIZE along the V dimension + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + STORE_FINAL_STATE: tl.constexpr, # whether to store final state + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + p_q = q + (bos * H + i_h) * K + tl.arange(0, BK) + p_k = k + (bos * H + i_h) * K + tl.arange(0, BK) + p_a = a + (bos * H + i_h) * K + tl.arange(0, BK) + p_b = b + (bos * H + i_h) * K + tl.arange(0, BK) + p_ha = ha + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_o = o + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + + mask_k = tl.arange(0, BK) < K + mask_v = (i_v * BV + tl.arange(0, BV)) < V + mask_h = mask_k[None, :] & mask_v[:, None] + + b_h = tl.zeros([BV, BK], dtype=tl.float32) + + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K * V + (tl.arange(0, BK)[None, :]) * V + ((i_v * BV + tl.arange(0, BV))[:, None]) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) + b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) + # to store + tmp = tl.sum(b_h * b_a[None, :], axis=1) + b_h += (tmp[:, None] * b_b[None, :] + b_k[None, :] * b_v[:, None]) + b_o = b_h * b_q[None, :] + b_o = tl.sum(b_o, axis=1) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + tl.store(p_ha, tmp.to(p_ha.dtype.element_ty), mask=mask_v) + p_q += K*H + p_k += K*H + p_o += V*H + p_v += V*H + p_ha += V*H + p_a += K*H + p_b += K*H + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K * V + (tl.arange(0, BK)[None, :]) * V + ((i_v * BV + tl.arange(0, BV))[:, None]) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_DHT': lambda args: args['dht'] is not None, + 'USE_DH0': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3] + ], + key=["BK", "BV"], + **autotune_cache_kwargs, +) +@triton.jit +def fused_recurrent_bwd_kernel( + # B: batch_size, H: n_heads, T: seq_len, D: b_dhead + # NV: number of split in the V dimension. NK: number of split in the K dimension + q, # query [B, H, L, K] + k, # key [B, H, L, V] + v, # value [B, H, L, V] + a, # a [B, H, L, K] + b, # b [B, H, L, K] + ha, # ha [B, H, L, V] + dht, # gradient of final state [B, H, K, V] + dh0, # gradient of initial state [B, H, K, V] + do, # gradient of output [B, H, L, V] + dq, # gradient of query [NV, B, H, L, K] + dk, # gradient of key [NV, B, H, L, K] + dv, # gradient of value [NK, B, H, L, V] + da, # gradient of a [NV, B, H, L, K] + db, # gradient of b [NV, B, H, L, K] + dha, # gradient of ha [NK, B, H, L, V] + h0, # initial state [B, H, K, V] + scale, # K ** -0.5 + cu_seqlens, # cu_seqlens + B, # batch_size + H, # n_heads + T, # seq_len + K: tl.constexpr, # K + V: tl.constexpr, # V + BK: tl.constexpr, # BLOCK SIZE along the K dimension + BV: tl.constexpr, # BLOCK SIZE along the V dimension + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state h0 + USE_DH0: tl.constexpr, # whether to use dh0 + USE_DHT: tl.constexpr, # whether to use dht + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + dk += i_v * B * H * K * T + db += i_v * B * H * K * T + dq += i_v * B * H * K * T + da += i_v * B * H * K * T + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + mask_k = tl.arange(0, BK) < K + mask_v = (tl.arange(0, BV) + i_v * BV) < V + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + i_v * BV + ha += (bos * H + i_h) * V + i_v * BV + a += (bos * H + i_h) * K + b += (bos * H + i_h) * K + do += (bos * H + i_h) * V + i_v * BV + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + i_v * BV + da += (bos * H + i_h) * K + db += (bos * H + i_h) * K + dha += (bos * H + i_h) * V + i_v * BV + + p_q = q + tl.arange(0, BK) + (T - 1) * H*K + p_k = k + tl.arange(0, BK) + (T - 1) * H*K + p_v = v + tl.arange(0, BV) + (T - 1) * H*V + p_ha = ha + tl.arange(0, BV) + (T - 1) * H*V + p_a = a + tl.arange(0, BK) + (T - 1) * H*K + p_b = b + tl.arange(0, BK) + (T - 1) * H*K + p_do = do + tl.arange(0, BV) + (T - 1) * H*V + p_dk = dk + tl.arange(0, BK) + (T - 1) * H*K + p_dv = dv + tl.arange(0, BV) + (T - 1) * H*V + p_dha = dha + tl.arange(0, BV) + (T - 1) * H*V + p_db = db + tl.arange(0, BK) + (T - 1) * H*K + p_da = da + tl.arange(0, BK) + (T - 1) * H*K + p_dq = dq + tl.arange(0, BK) + (T - 1) * H*K + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_DHT: + p_ht = dht + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + ((i_v * BV + tl.arange(0, BV))[None, :]) + b_dh += tl.load(p_ht, mask=mask_k[:, None] & mask_v[None, :], other=0).to(tl.float32) + + for _ in range(T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) + b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) + b_ha = tl.load(p_ha, mask=mask_v, other=0).to(tl.float32) + + b_dh += b_q[:, None] * b_do[None, :] + d_k = tl.sum(b_dh * b_v[None, :], axis=1) + d_v = tl.sum(b_dh * b_k[:, None], axis=0) + tl.store(p_dk, d_k.to(p_dk.dtype.element_ty), mask=mask_k) + tl.store(p_dv, d_v.to(p_dv.dtype.element_ty), mask=mask_v) + + b_dha = tl.sum(b_dh * b_b[:, None], axis=0) + tl.store(p_dha, b_dha.to(p_dha.dtype.element_ty), mask=mask_v) + b_db = tl.sum(b_dh * b_ha[None, :], axis=1) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=mask_k) + + b_dh += b_dha[None, :] * b_a[:, None] + p_do -= H*V + p_q -= H*K + p_k -= H*K + p_v -= H*V + p_dk -= H*K + p_dv -= H*V + p_b -= H*K + p_db -= H*K + p_a -= H*K + p_dha -= H*V + p_ha -= H*V + + if USE_DH0: + p_dh0 = dh0 + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_k[:, None] & mask_v[None, :]) + + tl.debug_barrier() + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + if USE_INITIAL_STATE: + mask_kv = mask_k[:, None] & mask_v[None, :] + p_h0 = h0 + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + ((i_v * BV + tl.arange(0, BV))[None, :]) + b_h += tl.load(p_h0, mask=mask_kv, other=0).to(tl.float32) + + p_k = k + tl.arange(0, BK) + p_v = v + tl.arange(0, BV) + p_ha = ha + tl.arange(0, BV) + p_do = do + tl.arange(0, BV) + p_dha = dha + tl.arange(0, BV) + p_da = da + tl.arange(0, BK) + p_dq = dq + tl.arange(0, BK) + p_b = b + tl.arange(0, BK) + + for i in range(0, T): + b_dha = tl.load(p_dha, mask=mask_v, other=0).to(tl.float32) + d_a = tl.sum(b_dha[None, :] * b_h, axis=1) + tl.store(p_da, d_a.to(p_da.dtype.element_ty), mask=mask_k) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) + b_ha = tl.load(p_ha, mask=mask_v, other=0).to(tl.float32) + b_h += b_k[:, None] * b_v[None, :] + b_b[:, None] * b_ha[None, :] + _d_q = b_h * b_do[None, :] + d_q = tl.sum(_d_q, axis=1) * scale + tl.store(p_dq, d_q.to(p_dq.dtype.element_ty), mask=mask_k) + + p_k += H*K + p_do += H*V + p_v += H*V + p_da += H*K + p_dha += H*V + p_ha += H*V + p_dq += H*K + p_b += H*K + + +class FusedRecurrentIPLRDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + BK = triton.next_power_of_2(K) + if output_final_state: + final_state = q.new_empty(B, H, K, V, dtype=torch.float32) + else: + final_state = None + + ha = torch.empty_like(v, dtype=torch.float32) + + def grid(meta): return ( + triton.cdiv(V, meta['BV']), + N * H, + ) + o = torch.empty_like(v) + fused_recurrent_fwd_kernel[grid]( + q=q, + k=k, + v=v, + a=a, + b=b, + o=o, + ha=ha, + h0=initial_state, + ht=final_state, + scale=scale, + cu_seqlens=cu_seqlens, + H=H, + T=T, + K=K, + V=V, + BK=BK, + ) + ctx.save_for_backward(q, k, v, a, b, ha, initial_state) + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + q, k, v, a, b, ha, initial_state = ctx.saved_tensors + B, T, H, K, V = *q.shape, v.shape[-1] + N = B if ctx.cu_seqlens is None else len(ctx.cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 64) + NV = triton.cdiv(V, BV) + scale = ctx.scale + + dq = q.new_empty(NV, *q.shape) + dk = k.new_empty(NV, *k.shape) + da = a.new_empty(NV, *a.shape) + db = b.new_empty(NV, *b.shape) + dv = torch.empty_like(v) + dha = torch.empty_like(ha) + grid = (NV, N * H) + + if initial_state is not None and initial_state.requires_grad: + dh0 = torch.empty_like(initial_state, dtype=torch.float32) + else: + dh0 = None + + fused_recurrent_bwd_kernel[grid]( + q=q, + k=k, + v=v, + a=a, + b=b, + ha=ha, + dht=dht, + dh0=dh0, + do=do, + dq=dq, + dk=dk, + dv=dv, + da=da, + db=db, + dha=dha, + h0=initial_state, + scale=scale, + cu_seqlens=ctx.cu_seqlens, + B=B, + H=H, + T=T, + K=K, + V=V, + BK=BK, + BV=BV, + ) + dq = dq.sum(0) + dk = dk.sum(0) + da = da.sum(0) + db = db.sum(0) + return dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), None, dh0, None, None + + +def fused_recurrent_iplr_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. + + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]` + k (torch.Tensor): + keys of shape `[B, T, H, K]` + v (torch.Tensor): + values of shape `[B, T, H, V]` + a (torch.Tensor): + as of shape `[B, T, H, K]` + b (torch.Tensor): + bs of shape `[B, T, H, K]` + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[B, H, K, V]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = q.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + o, final_state = FusedRecurrentIPLRDeltaRuleFunction.apply( + q, + k, + v, + a, + b, + scale, + initial_state, + output_final_state, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/generalized_delta_rule/iplr/naive.py b/fla/ops/generalized_delta_rule/iplr/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..2368b2baa0a2357bedc146f84c97b785cbaf6506 --- /dev/null +++ b/fla/ops/generalized_delta_rule/iplr/naive.py @@ -0,0 +1,68 @@ + +import torch +from einops import rearrange + + +# S_t = S_t @ (I + alpha_t beta_t^T) + v_t k_t^T +# q, k, alpha, beta [B, H, L, D_K] +# v [B, H, L, D_V] +def iplr_recurrence(q, k, v, alpha, beta, initial_state=None, output_final_state=True): + orig_dtype = q.dtype + b, h, l, d_k = q.shape + q, k, v, beta = map(lambda x: x.float(), [q, k, v, beta]) + d_v = v.shape[-1] + o = torch.zeros_like(v) + S = torch.zeros(b, h, d_k, d_v).to(v) + q = q * (d_k ** -0.5) + + if initial_state is not None: + S += initial_state + + for i in range(l): + _k = k[:, :, i] + _q = q[:, :, i] + _v = v[:, :, i] + _alpha = alpha[:, :, i] + _beta = beta[:, :, i] + _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _alpha[..., None]).sum(-2, keepdim=True) * _beta[..., None] + S = S + _kv + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) + S = None if output_final_state is False else S + return o.to(orig_dtype), S + + +def iplr_chunkwise(q, k, v, alpha, beta, initial_state=None, output_final_state=True, chunk_size=32): + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * (d_k ** -0.5) + v = v + assert l % chunk_size == 0 + + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S += initial_state + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, alpha, beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, alpha, beta]) + + v2 = (alpha @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v + attn = (alpha @ beta.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2) + + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + u = attn @ v2 + w = attn @ alpha + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i, u_i, w_i, beta_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], beta[:, :, i] + o_1 = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v_i + v2_i = u_i + w_i @ S + o_2 = (q_i @ beta_i.transpose(-1, -2)).masked_fill_(mask, 0) @ (v2_i) + o_3 = q_i @ S + o[:, :, i] = o_1 + o_2 + o_3 + S = S + k_i.transpose(-1, -2) @ v_i + beta_i.transpose(-1, -2) @ v2_i + S = None if output_final_state is False else S + return rearrange(o, 'b h n c d -> b h (n c) d'), S diff --git a/fla/ops/generalized_delta_rule/iplr/wy_fast.py b/fla/ops/generalized_delta_rule/iplr/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..6a7965f7c1c4903f234d0c55a79578448ab7c31e --- /dev/null +++ b/fla/ops/generalized_delta_rule/iplr/wy_fast.py @@ -0,0 +1,304 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + ], + key=['BK'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_fwd_kernel_chunk32( + a, + b, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BC: tl.constexpr, # dummy placeholder + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_b = tl.load(p_b, boundary_check=(0, 1)) + b_A += tl.dot(b_a, b_b) + + b_A = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A, 0) + for i in range(1, BT): + mask = tl.arange(0, BT) == i + b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BT) < i) + b_A = tl.where(mask[:, None], b_a, b_A) + b_A += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + ], + key=['BK'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_fwd_kernel_chunk64( + a, + b, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_A2 = tl.zeros([BC, BC], dtype=tl.float32) + b_A3 = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_a1 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_a2 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + BC, i_k * BK), (BC, BK), (1, 0)) + p_b1 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BC), (0, 1)) + p_b2 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT + BC), (BK, BC), (0, 1)) + b_a1 = tl.load(p_a1, boundary_check=(0, 1)) + b_a2 = tl.load(p_a2, boundary_check=(0, 1)) + b_b1 = tl.load(p_b1, boundary_check=(0, 1)) + b_b2 = tl.load(p_b2, boundary_check=(0, 1)) + b_A += tl.dot(b_a1, b_b1, allow_tf32=False) + b_A2 += tl.dot(b_a2, b_b2, allow_tf32=False) + b_A3 += tl.dot(b_a2, b_b1, allow_tf32=False) + + b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) + b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) + + for i in range(1, BC): + mask = tl.arange(0, BC) == i + b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) + b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BC) < i) + b_a2 = b_a2 + tl.sum(b_a2[:, None] * b_A2, 0) * (tl.arange(0, BC) < i) + b_A = tl.where(mask[:, None], b_a, b_A) + b_A2 = tl.where(mask[:, None], b_a2, b_A2) + + # blockwise computation of lower triangular matrix's inverse + # i.e., [A11, 0; A21, A22]^-1 = [A11^-1, 0; -A22^-1 A21 A11^-1, A22^-1] + b_A += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] + b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] + b_A3 = tl.dot(tl.dot(b_A2, b_A3, allow_tf32=False), b_A, allow_tf32=False) + + p_A1 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) + p_A2 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) + p_A3 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) + p_A4 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) + tl.store(p_A1, b_A.to(p_A1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A3, b_A3.to(p_A3.dtype.element_ty), boundary_check=(0, 1)) + # causal mask + tl.store(p_A4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A4.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def wu_fwd_kernel( + w, + u, + a, + k, + v, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_Aak = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_a = tl.load(p_a, boundary_check=(0, 1)) + b_w = tl.dot(b_A, b_a) + b_Aak += tl.dot(b_a, tl.trans(b_k)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + b_Aak = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_Aak, 0) + b_Aak = b_Aak.to(k.dtype.element_ty) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.dot(b_Aak, b_v).to(v.dtype.element_ty) + b_u = tl.dot(b_A, b_v) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + +def prepare_wy_repr_fwd( + a: torch.Tensor, + b: torch.Tensor, + v: torch.Tensor, + k: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K = a.shape + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BC = min(BT, 32) + BK = min(max(triton.next_power_of_2(K), 16), 64) + + A = torch.empty(B, T, H, BT, device=a.device, dtype=a.dtype) + fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 + + fwd_fn[(NT, B * H)]( + a=a, + b=b, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BK=BK, + BC=BC, + ) + w, u = wu_fwd( + a=a, + v=v, + k=k, + A=A, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return w, u, A + + +def wu_fwd( + a: torch.Tensor, + v: torch.Tensor, + k: torch.Tensor, + A: torch.Tensor, + cu_seqlens: torch.LongTensor | None, + chunk_size: int, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *a.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + u = torch.empty_like(v) + w = torch.empty_like(a) + wu_fwd_kernel[(NT, B*H)]( + a=a, + v=v, + w=w, + u=u, + A=A, + k=k, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +fwd_prepare_wy_repr = prepare_wy_repr_fwd + +fwd_wu = wu_fwd diff --git a/fla/ops/gla/__init__.py b/fla/ops/gla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa786e039f493817fb353072917587534c863a8d --- /dev/null +++ b/fla/ops/gla/__init__.py @@ -0,0 +1,10 @@ + +from .chunk import chunk_gla +from .fused_chunk import fused_chunk_gla +from .fused_recurrent import fused_recurrent_gla + +__all__ = [ + 'chunk_gla', + 'fused_chunk_gla', + 'fused_recurrent_gla', +] diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..adcf2350ad85205da3f3d3cbf5a1237bab3331bf --- /dev/null +++ b/fla/ops/gla/chunk.py @@ -0,0 +1,1376 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cumsum import chunk_local_cumsum +from fla.ops.utils.op import exp, exp2 +from fla.utils import autotune_cache_kwargs, check_shared_mem, input_guard + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_fwd_A_kernel_intra_sub_inter( + q, + k, + g, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gn = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_qg = b_q * exp(b_g - b_gn[None, :]) * scale + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[:, None] - b_gk) + + # [BC, BC] using tf32 to improve precision here. + b_A += tl.dot(b_qg, b_kg) + + p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3] + ], + key=["BK", "BT"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_fwd_A_kernel_intra_sub_intra( + q, + k, + g, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_j = i_i + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + o_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_j * BC + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_k = k + (i_t * BT + i_j * BC) * H*K + o_k + p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_A = tl.sum(b_q * b_k[None, :] * exp(b_g - b_gk[None, :]), 1) * scale + + tl.store(A + o_A + j, b_A, mask=m_A) + p_k += H*K + p_gk += H*K + + tl.debug_barrier() + b_A = tl.zeros([BC, BC], dtype=tl.float32) + tl.store(A + o_A[:, None] + o_i, b_A, mask=m_A[:, None] & (o_i[:, None] < o_i)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BC', 'BK'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_fwd_A_kernel_intra_sub_intra_split( + q, + k, + g, + A, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = i_tc // NC, i_tc % NC + i_j = i_i + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = i_k * BK + tl.arange(0, BK) + o_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BC + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += ((i_k * all + bos) * H + i_h) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_k = k + (i_t * BT + i_j * BC) * H*K + o_k + p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_A = tl.sum(b_q * b_k[None, :] * exp(b_g - b_gk[None, :]), 1) * scale + tl.store(A + o_A + j, b_A, mask=m_A) + p_k += H*K + p_gk += H*K + + tl.debug_barrier() + b_A = tl.zeros([BC, BC], dtype=tl.float32) + tl.store(A + o_A[:, None] + o_i, b_A, mask=m_A[:, None] & (o_i[:, None] < o_i)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + ], + key=['BC'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( + A, + A2, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + NK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + if i_t * BT + i_c * BC >= T: + return + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(0, NK): + p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) + b_A += tl.load(p_A, boundary_check=(0, 1)) + p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) + tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT', 'TRANSPOSE_STATE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_EXP2: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t.to(tl.int64) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + # [BT, BK] + if USE_EXP2: + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + else: + b_qg = (b_q * exp(b_g)).to(b_q.dtype) + b_h = tl.load(p_h, boundary_check=(0, 1)) + if i_k >= 0: + if TRANSPOSE_STATE: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + else: + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + b_o *= scale + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.).to(b_v.dtype) + b_o += tl.dot(b_A, b_v) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_bwd_kernel_intra( + q, + k, + g, + dA, + dq, + dk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_k, i_i = i_kc // NC, i_kc % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + if i_t * BT + i_i * BC >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_g = tl.make_block_ptr(g + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + + b_dq = tl.zeros([BC, BK], dtype=tl.float32) + if i_i > 0: + p_gn = g + (bos + i_t * BT + i_i * BC) * H*K + i_h*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[None, :] - b_gk) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + + b_dq += tl.dot(b_dA, b_kg) + b_dq *= exp(b_g - b_gn[None, :]) + + o_i = tl.arange(0, BC) + m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC + p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gkj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) + # [BK,] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) + p_kj += H*K + p_gkj += H*K + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + # [BC, BK] + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + i_h * K + o_k + + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(i_i + 1, NC): + p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_gq = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_gq = tl.load(p_gq, boundary_check=(0, 1)) + b_qg = b_q * tl.where(m_j[:, None], exp(b_gq - b_gn[None, :]), 0) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dk += tl.dot(b_dA, b_qg) + b_dk *= exp(b_gn[None, :] - b_g) + o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) + p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gqj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dA = tl.load(dA + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_gqj = tl.load(p_gqj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp(b_gqj[None, :] - b_g), 0.) + p_qj += H*K + p_gqj += H*K + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BV', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_bwd_kernel_dA( + v, + do, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + b_dA += tl.dot(b_do, b_v) + + p_dA = tl.make_block_ptr(dA + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + b_dA = tl.where(m_s, b_dA * scale, 0.) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_bwd_kernel_dv( + k, + g, + A, + do, + dh, + dv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_do = tl.make_block_ptr(do + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A, 0.) + # (SY 09/17) important to disallow tf32 here to maintain a good precision. + b_dv = tl.dot(b_A, b_do.to(b_A.dtype)) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_gk = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_gn = g + (bos + min(i_t * BT + BT, T) - 1)*H*K + i_h * K + o_k + p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + b_gn = exp(tl.load(p_gn, mask=m_k, other=0)[None, :] - b_gk) + b_k = (b_k * b_gn).to(b_k.dtype) + # [BT, BV] + # (SY 09/17) it is ok to have bf16 interchunk gradient contribution here + b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in [2, 4, 8] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gla_bwd_kernel_inter( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dq2, + dk2, + dg, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + g += (bos * H + i_h) * K + h += (i_tg * H + i_h) * K*V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K*V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dq2 += (bos * H + i_h) * K + dk2 += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0) + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + # [BK] + b_dgk += tl.sum(b_h * b_dh, axis=0) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + + b_dgk *= exp(b_gn) + b_dq *= scale + b_dq = b_dq * exp(b_gk) + b_dk = b_dk * exp(b_gn[None, :] - b_gk) + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dgk += tl.sum(b_dk * b_k, axis=0) + b_dq += tl.load(p_dq, boundary_check=(0, 1)) + b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dg = b_q * b_dq - b_k * b_dk + # tl.debug_barrier() + b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] + # Buggy due to strange triton compiler issue. + # m_s = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], 1., 0.) + # b_dg = tl.dot(m_s, b_dg) + b_dgk[None, :] + p_dq = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_intra_gk( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = k.shape + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BC = min(16, BT) + NC = triton.cdiv(BT, BC) + + A = q.new_empty(B, T, H, BT, dtype=torch.float) + grid = (NT, NC * NC, B * H) + chunk_gla_fwd_A_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=g, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + # load the entire [BC, K] blocks into SRAM at once + if K <= 256: + BK = max(triton.next_power_of_2(K), 16) + chunk_gla_fwd_A_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=g, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + # split then merge + else: + BK = min(128, triton.next_power_of_2(K)) + NK = triton.cdiv(K, BK) + A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) + + grid = (NK, NT * NC, B * H) + chunk_gla_fwd_A_kernel_intra_sub_intra_split[grid]( + q=q, + k=k, + g=g, + A=A_intra, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_gla_fwd_A_kernel_intra_sub_intra_merge[grid]( + A=A_intra, + A2=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + BC=BC, + NK=NK, + ) + return A + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + use_exp2: bool = False, + transpose_state_layout: bool = False, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + # Please ensure zeros, since vllm will use padding v + o = torch.zeros_like(v) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + USE_EXP2=use_exp2, + TRANSPOSE_STATE=transpose_state_layout, + ) + return o + + +def chunk_gla_bwd_dA( + v: torch.Tensor, + do: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, V = v.shape + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BV = min(64, triton.next_power_of_2(V)) + + dA = v.new_empty(B, T, H, BT, dtype=torch.float) + grid = (NT, B * H) + chunk_gla_bwd_kernel_dA[grid]( + v=v, + do=do, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + V=V, + BT=BT, + BV=BV, + ) + return dA + + +def chunk_gla_bwd_dv( + k: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dv = torch.empty_like(do) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + chunk_gla_bwd_kernel_dv[grid]( + k=k, + g=g, + A=A, + do=do, + dh=dh, + dv=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dv + + +def chunk_gla_bwd_dqk_intra( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + dA: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = q.shape + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + grid = (NK * NC, NT, B * H) + chunk_gla_bwd_kernel_intra[grid]( + q=q, + k=k, + g=g, + dA=dA, + dq=dq, + dk=dk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + return dq, dk + + +def chunk_gla_bwd_dqkg( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dg = torch.empty_like(g) + dq2 = torch.empty_like(dq) + dk2 = torch.empty_like(dk) + def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) + chunk_gla_bwd_kernel_inter[grid]( + q=q, + k=k, + v=v, + g=g, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dq2=dq2, + dk2=dk2, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dq2, dk2, dg + + +def chunk_gla_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_cumsum: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if g_cumsum is None: + g_cumsum = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens) + + h, ht = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=g_cumsum, + gv=None, + h0=initial_state, + output_final_state=output_final_state, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + + # the intra A is kept in fp32 + # the computation has very marginal effect on the entire throughput + A = chunk_gla_fwd_intra_gk( + q=q, + k=k, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + o = chunk_gla_fwd_o_gk( + q=q, + v=v, + g=g_cumsum, + A=A, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return g_cumsum, A, h, ht, o + + +def chunk_gla_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_cumsum: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + h: torch.Tensor, + A: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + if g_cumsum is None: + g_cumsum = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens) + + if h is None: + h, _ = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=g_cumsum, + gv=None, + h0=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=None, + gk=g_cumsum, + gv=None, + do=do, + h0=initial_state, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + ) + + dv = chunk_gla_bwd_dv( + k=k, + g=g_cumsum, + A=A, + do=do, + dh=dh, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # dq dk in fp32 + dA = chunk_gla_bwd_dA( + v=v, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dq, dk = chunk_gla_bwd_dqk_intra( + q=q, + k=k, + g=g_cumsum, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dq, dk, dg = chunk_gla_bwd_dqkg( + q=q, + k=k, + v=v, + h=h, + g=g_cumsum, + do=do, + dh=dh, + dq=dq, + dk=dk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, dg, dh0 + + +class ChunkGLAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q, + k, + v, + g, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ): + chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g_cumsum, A, _, ht, o = chunk_gla_fwd( + q=q, + k=k, + v=v, + g=g, + g_cumsum=None, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + # recompute g_cumsum in bwd pass + if g.dtype != torch.float: + g_cumsum = None + else: + g = None + ctx.save_for_backward(q, k, v, g, g_cumsum, initial_state, A, chunk_indices) + ctx.chunk_size = chunk_size + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o, ht + + @staticmethod + @input_guard + def backward(ctx, do, dht): + q, k, v, g, g_cumsum, initial_state, A, chunk_indices = ctx.saved_tensors + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens + dq, dk, dv, dg, dh0 = chunk_gla_bwd( + q=q, + k=k, + v=v, + g=g, + g_cumsum=g_cumsum, + scale=scale, + h=None, + A=A, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), dg, None, dh0, None, None, None + + +@torch.compiler.disable +def chunk_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: int | None = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H, K]`. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gla import chunk_gla + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = chunk_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = q.shape[-1] ** -0.5 + if initial_state is not None: + assert initial_state.dtype == torch.float32, "initial_state must be in float32." + assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." + assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + o, final_state = ChunkGLAFunction.apply(q, k, v, g, scale, initial_state, output_final_state, cu_seqlens, cu_seqlens_cpu) + return o, final_state diff --git a/fla/ops/gla/fused_chunk.py b/fla/ops/gla/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..6cbf86fe9c7b5ebe6d37df528642c4c7a23e3e21 --- /dev/null +++ b/fla/ops/gla/fused_chunk.py @@ -0,0 +1,5 @@ + +def fused_chunk_gla( + **kwargs, +): + raise NotImplementedError("`fused_chunk_gla` is deprecated. Please use `chunk_gla` instead.") diff --git a/fla/ops/gla/fused_recurrent.py b/fla/ops/gla/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..5f7c99a7633d0f7725b9fab6d6255ab83bac0f66 --- /dev/null +++ b/fla/ops/gla/fused_recurrent.py @@ -0,0 +1,109 @@ +# Copyright (c) 2024, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.common.fused_recurrent import fused_recurrent + + +def fused_recurrent_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + scale: int | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + gk (torch.Tensor): + Forget gates of shape `[B, T, H, K]`. + gv (torch.Tensor): + Forget gates of shape `[B, T, H, V]` applied to values. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gla import fused_recurrent_gla + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_recurrent_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = fused_recurrent( + q=q, + k=k, + v=v, + g=None, + gk=gk, + gv=gv, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/gla/naive.py b/fla/ops/gla/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..2fcbd5089d5178ad2cb3e26e2c75cf9b7d8165d7 --- /dev/null +++ b/fla/ops/gla/naive.py @@ -0,0 +1,39 @@ + + +import torch + + +def ceildiv(a, b): + return -(a // -b) + + +def naive_recurrent_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gk: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, +): + dtype = q.dtype + q, k, v, gk = map(lambda x: x.transpose(1, 2).float(), (q, k, v, gk)) + B, H, T, K, V = *q.shape, v.shape[-1] + o = torch.zeros_like(v) + scale = K ** -0.5 + + h = q.new_zeros(B, H, K, V, dtype=torch.float32) + if initial_state is not None: + h += initial_state.float() + + for i in range(T): + q_i = q[:, :, i] * scale + k_i = k[:, :, i] + v_i = v[:, :, i] + gk_i = gk[:, :, i].exp() + kv_i = k_i[..., None] * v_i[..., None, :] + h = h * gk_i[..., None] + kv_i + o[:, :, i] = (q_i[..., None] * h).sum(-2) + + if not output_final_state: + h = None + return o.transpose(1, 2).to(dtype), h diff --git a/fla/ops/gsa/__init__.py b/fla/ops/gsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1d827b221e7a0ac104e57a506c7576e7f871a19a --- /dev/null +++ b/fla/ops/gsa/__init__.py @@ -0,0 +1,8 @@ + +from .chunk import chunk_gsa +from .fused_recurrent import fused_recurrent_gsa + +__all__ = [ + 'chunk_gsa', + 'fused_recurrent_gsa', +] diff --git a/fla/ops/gsa/chunk.py b/fla/ops/gsa/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..7987d013788000753228c044dcaff3ee017a28fe --- /dev/null +++ b/fla/ops/gsa/chunk.py @@ -0,0 +1,1155 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl +from einops import reduce + +from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from fla.ops.gla.chunk import chunk_gla_bwd, chunk_gla_fwd +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cumsum import chunk_local_cumsum +from fla.ops.utils.op import exp +from fla.ops.utils.softmax import softmax_bwd, softmax_fwd +from fla.utils import autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_k_kernel_inter( + q, + k, + h, + g, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BT] + b_A += tl.dot(b_q, b_k) + p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o = b_o * exp(b_g) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + # [BT, BT] + b_A = tl.where(m_s, b_A, 0.) + if i_v == 0: + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_k_kernel_intra( + v, + g, + o, + A, + cu_seqlens, + chunk_indices, + T, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = g + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_o = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(0, i_i): + p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[None, :] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_o += tl.dot(b_A, b_vg) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o *= exp(b_g - b_gn[None, :]) + + o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + # [BC,] + b_A = tl.load(A + o_A + j, mask=m_A, other=0) + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC, BV] + b_vg = b_v[None, :] * exp(b_g - b_gv[None, :]) + # avoid 0 * inf = inf + b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) + p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8] + ], + key=["BT"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_dA( + v, + g, + do, + dA, + chunk_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) + + # [BC, BC] + b_dA = tl.zeros([BC, BC], dtype=tl.float32) + if i_i > i_j: + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gn = g + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0.) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g - b_gn[None, :])).to(b_do.dtype) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[:, None] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_dA = tl.dot(b_do, b_vg) * scale + elif i_i == i_j: + p_g = tl.make_block_ptr(g + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + m_v = o_v < V + + o_i = tl.arange(0, BC) + # [BC, BC] + m_dA = o_i[:, None] >= o_i[None, :] + for j in range(0, min(BC, T - i_t * BT - i_j * BC)): + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC,] + b_dAj = tl.sum(b_do * b_v[None, :] * exp(b_g - b_gv[None, :]), 1) + b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) + + p_v += H*V + p_gv += H*V + b_dA = tl.where(m_dA, b_dA, 0.) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_dqkvg( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_i = tl.arange(0, BT) + o_t = min(i_t * BT + BT, T) + m_s = o_i[:, None] >= o_i[None, :] + + p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + m_v = o_v < V + + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_gv = exp(b_gn[None, :] - b_g) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g)).to(b_do.dtype) + # [BK, BV] + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BV] + b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp(b_gn) + + b_dh = b_dh.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) * scale + b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) + # [BT, BV] + b_dv = tl.dot(b_k, b_dh) * b_gv + # [BV] + b_dg += tl.sum(b_dv * b_v, 0) + + if i_k == 0: + b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] + else: + b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] + + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BT] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BT, BK] + b_dq += tl.dot(b_dA, b_k) + b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_intra_dvg( + v, + g, + o, + A, + do, + dv, + dg, + cu_seqlens, + chunk_indices, + T, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (BT, T), (1, HQ*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) + + m_j = (i_t * BT + i_j * BC + o_i) < T + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * tl.where(m_j[:, None], exp(b_g - b_gn[None, :]), 0) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + # [BC, BV] + b_dv += tl.dot(b_A, b_do.to(b_A.dtype)) + b_dv *= exp(b_gn[None, :] - b_gv) + + p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v + p_A = A + (bos + i_t*BT + i_i*BC) * HQ*BT + i_hq * BT + i_i * BC + o_i + p_do = do + (bos + i_t*BT + i_i*BC) * HQ*V + i_hq * V + o_v + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_A = tl.load(p_A) + # [BV,] + b_g = tl.load(p_g, mask=m_v, other=0) + b_do = tl.load(p_do, mask=m_v, other=0) + # [BC, BV] + m_i = o_i[:, None] <= j + b_dv += tl.where(m_i, exp(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) + + p_g += H * V + p_A += HQ * BT + p_do += HQ * V + p_o = tl.make_block_ptr(o + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + + b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) + b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) + b_dg = b_o * b_do - b_v * b_dv + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gsa_fwd_v( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float = 1., + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + _, A, h, ht, o = chunk_gla_fwd( + q=q, + k=k, + v=v, + g=None, + g_cumsum=g, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return A, h, ht, o + + +def chunk_gsa_fwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h0: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + HQ = q.shape[2] + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NG = HQ // H + + h, ht = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=False, + ) + o = v.new_empty(B, T, HQ, V) + A = q.new_empty(B, T, HQ, BT) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) + chunk_gsa_fwd_k_kernel_inter[grid]( + q, + k, + h, + g, + o, + A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + ) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + chunk_gsa_fwd_k_kernel_intra[grid]( + v, + g, + o, + A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + num_warps=4, + num_stages=2, + ) + return A, h, ht, o + + +def chunk_gsa_bwd_v( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h0: torch.Tensor, + h: torch.Tensor, + A: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dg: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + dq, dk, dv, dg, dh0 = chunk_gla_bwd( + q=q, + k=k, + v=v, + g=None, + g_cumsum=g, + scale=scale, + initial_state=h0, + h=h, + A=A, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, dg, dh0 + + +def chunk_gsa_bwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + h0: torch.Tensor, + o: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dg: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + HQ = q.shape[2] + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + NG = HQ // H + + if h is None: + h, _ = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=False, + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=None, + gk=None, + gv=g, + do=do, + h0=h0, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=True, + ) + dA = q.new_empty(NV, B, T, HQ, BT) + grid = (NV, NT * NC * NC, B * HQ) + chunk_gsa_bwd_k_kernel_dA[grid]( + v, + g, + do, + dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + ) + dA = dA.sum(0, dtype=dA.dtype) + + A = do.new_empty(NK, B, T, HQ, BT) + dq = torch.empty_like(q) + dk = k.new_empty(B, T, HQ, K) + dv = v.new_empty(NK, B, T, HQ, V) + dgv = g.new_empty(NK, B, T, HQ, V, dtype=torch.float) + grid = (NK, NT, B * HQ) + chunk_gsa_bwd_k_kernel_dqkvg[grid]( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + NG=NG, + ) + A = A.sum(0, dtype=A.dtype) + dv = dv.sum(0, dtype=dv.dtype) + dgv = dgv.sum(0, dtype=dgv.dtype) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + chunk_gsa_bwd_k_kernel_intra_dvg[grid]( + v, + g, + o, + A, + do, + dv, + dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + num_warps=4, + num_stages=2, + ) + dg = dgv.add_(chunk_local_cumsum(dg, chunk_size=BT, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)) + + return dq, dk, dv, dg, dh0 + + +def chunk_gsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + Ak, hk, hkt, ok = chunk_gsa_fwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + output_final_state=output_final_state, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # p is kept in fp32 for safe softmax backward + p = softmax_fwd(ok, dtype=torch.float) + + qv = p.to(q.dtype) + Av, hv, hvt, ov = chunk_gsa_fwd_v( + q=qv, + k=s, + v=v, + g=g, + scale=1., + initial_state=hv0, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return Ak, hk, hkt, ok, p, Av, hv, hvt, ov + + +def chunk_gsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + ok: torch.Tensor, + p: torch.Tensor, + A: tuple[torch.Tensor, torch.Tensor], + h: tuple[torch.Tensor, torch.Tensor], + initial_state: tuple[torch.Tensor, torch.Tensor] | None, + scale: float, + do: torch.Tensor, + dht: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + + _, Av = A + hk, hv = h + dhkt, dhvt = dht + + qv = p.to(q.dtype) + dqv, dsv, dv, dg, dhv0 = chunk_gsa_bwd_v( + q=qv, + k=s, + v=v, + g=g, + h0=hv0, + h=hv, + A=Av, + do=do, + dht=dhvt, + dg=None, + scale=1., + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # softmax gradient, equivalent to: + # dok = qv * (dqv - (qv * dqv).sum(-1, True)) + dok = softmax_bwd(p, dqv, dtype=ok.dtype) + + dq, dk, dsk, dg, dhk0 = chunk_gsa_bwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + h=hk, + o=ok, + do=dok, + dht=dhkt, + dg=dg, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + ds = dsv.add_(dsk) + if q.shape[1] != k.shape[1]: + dk, dv, ds, dg = map(lambda x: reduce(x, 'b (h g) ... -> b h ...', 'sum', h=k.shape[1]), (dk, dv, ds, dg)) + dg = dg.to(s.dtype) + return dq, dk, dv, ds, dg, dhk0, dhv0 + + +class ChunkGSAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + scale: float, + hk0: torch.Tensor | None, + hv0: torch.Tensor | None, + output_final_state: bool, + checkpoint_level: int, + cu_seqlens: torch.LongTensor | None, + cu_seqlens_cpu: torch.LongTensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g_org, g = g, chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + Ak, hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + if checkpoint_level >= 1: + del g + g = g_org + if checkpoint_level > 1: + del hk + del hv + hk, hv = None, None + else: + hk0, hv0 = None, None + + ctx.save_for_backward(q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv, chunk_indices) + ctx.checkpoint_level = checkpoint_level + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + ctx.chunk_size = chunk_size + return ov, hkt, hvt + + @staticmethod + @input_guard + def backward(ctx, dov, dhkt=None, dhvt=None): + q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv, chunk_indices = ctx.saved_tensors + scale = ctx.scale + cu_seqlens = ctx.cu_seqlens + chunk_size = ctx.chunk_size + + if ctx.checkpoint_level >= 1: + g = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + dq, dk, dv, ds, dg, dhk0, dhv0 = chunk_gsa_bwd( + q=q, + k=k, + v=v, + s=s, + g=g, + ok=ok, + p=p, + A=(None, Av), + h=(hk, hv), + initial_state=(hk0, hv0), + scale=scale, + do=dov, + dht=(dhkt, dhvt), + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, ds, dg, None, dhk0, dhv0, None, None, None, None + + +@torch.compiler.disable +def chunk_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor | None = None, + scale: int | None = None, + initial_state: tuple[torch.Tensor] | None = None, + output_final_state: bool | None = False, + checkpoint_level: int | None = 2, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool | None = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]`.. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + GQA is performed if `H` is not equal to `HQ`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + s (torch.Tensor): + slot representations of shape `[B, T, H, M]`.. + g (torch.Tensor): + Forget gates of shape `[B, T, H, M]` applied to keys. + If not provided, this function is equivalent to vanilla ABC. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[Tuple[torch.Tensor]]): + Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state tuple, having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. + Default: `False`. + checkpoint_level (Optional[int]): + Checkpointing level; higher values will save more memories and do more recomputations during backward. + Default: `2`: + - Level `0`: no memory saved, no recomputation. + - Level `1`: recompute the fp32 cumulative values during backward. + - Level `2`: recompute the fp32 cumulative values and forward hidden states during backward. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (Tuple[torch.Tensor]): + Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` if `output_final_state=True`. + `None` otherwise. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gsa import fused_recurrent_gsa + # inputs with equal lengths + >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> s = torch.randn(B, T, H, M, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) + >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) + >>> o, (hk, hv) = chunk_gsa( + q, k, v, s, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, (hk_var, hv_var) = chunk_gsa( + q, k, v, s, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert hk.allclose(hk_var) + >>> assert hv.allclose(hv_var) + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.", + ) + assert checkpoint_level in [0, 1, 2] + if g is None: + # TODO: this 3 steps took huge amount of time, ought to be optimized + z = s.float().logcumsumexp(2) + g = torch.cat((z[:, :, :1], z[:, :, :-1]), 1) - z + s = torch.exp(s - z).to(k.dtype) + if scale is None: + scale = q.shape[-1] ** -0.5 + + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + o, *final_state = ChunkGSAFunction.apply( + q, + k, + v, + s, + g, + scale, + hk0, + hv0, + output_final_state, + checkpoint_level, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/gsa/fused_recurrent.py b/fla/ops/gsa/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..82f49a421b66a6db0ea74bd4b2699b6014fe124e --- /dev/null +++ b/fla/ops/gsa/fused_recurrent.py @@ -0,0 +1,534 @@ +# Copyright (c) 2024, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.common.fused_recurrent import fused_recurrent_bwd_kernel, fused_recurrent_fwd_kernel +from fla.ops.utils.op import exp +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.jit +def fused_recurrent_gsa_inference_kernel( + q, + k, + v, + s, + g, + o, + hk0, + hv0, + hkt, + hvt, + scale, + K: tl.constexpr, + V: tl.constexpr, + M: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, +): + i_bh = tl.program_id(0) + i_bg = i_bh // NG + + b_s = tl.load(s + i_bg * M + tl.arange(0, M)).to(tl.float32) + b_g = tl.load(g + i_bg * M + tl.arange(0, M)).to(tl.float32) + b_g = exp(b_g) + + b_ok = tl.zeros([M], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + + p_hk0 = hk0 + i_bg * K * M + (o_k[None, :]) * M + tl.arange(0, M)[:, None] + # [BK,] + mask_k = o_k < K + # [M, BK] + mask_hk = (tl.arange(0, M) < M)[:, None] & mask_k[None, :] + # [M, BK] + b_hk = tl.load(p_hk0, mask=mask_hk, other=0.).to(tl.float32) + # [BK,] + b_q = tl.load(q + i_bh * K + o_k, mask=mask_k, other=0.).to(tl.float32) * scale + b_k = tl.load(k + i_bg * K + o_k, mask=mask_k, other=0.).to(tl.float32) + b_hk = b_hk * b_g[:, None] + b_k[None, :] * b_s[:, None] + b_ok += tl.sum(b_hk * b_q[None, :], axis=1) + + if i_bh % NG == 0: + p_hkt = hkt + i_bg * K * M + o_k[None, :] * M + tl.arange(0, M)[:, None] + tl.store(p_hkt, b_hk.to(p_hkt.dtype.element_ty), mask=mask_hk) + + b_qv = tl.softmax(b_ok) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + + p_hv0 = hv0 + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] + # [BV,] + mask_v = o_v < V + # [BV, M] + mask_hv = mask_v[:, None] & (tl.arange(0, M) < M)[None, :] + # [BV, M] + b_hv = tl.load(p_hv0, mask=mask_hv, other=0).to(tl.float32) + # [BV,] + b_v = tl.load(v + i_bg * V + o_v, mask=mask_v, other=0).to(tl.float32) + b_hv = b_hv * b_g[None, :] + b_s[None, :] * b_v[:, None] + b_ov = tl.sum(b_hv * b_qv[None, :], axis=1) + + tl.store(o + i_bh * V + o_v, b_ov.to(o.dtype.element_ty), mask=mask_v) + + if i_bh % NG == 0: + p_hvt = hvt + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] + tl.store(p_hvt, b_hv.to(p_hvt.dtype.element_ty), mask=mask_hv) + + +def fused_recurrent_gsa_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + scale: float = 1., +) -> torch.Tensor: + B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + HQ = q.shape[2] + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NG = HQ // H + + if initial_state != (None, None) and initial_state is not None: + hk0, hv0 = initial_state + else: + hk0, hv0 = q.new_zeros(B, H, K, M, dtype=torch.float), q.new_zeros(B, H, M, V, dtype=torch.float) + + hkt, hvt = None, None + if output_final_state: + if NG == 1: + hkt, hvt = hk0, hv0 + else: + hkt, hvt = q.new_empty(B, H, K, M, dtype=torch.float), q.new_empty(B, H, M, V, dtype=torch.float) + + o = v.new_empty(B, T, HQ, V) + grid = (B * HQ,) + fused_recurrent_gsa_inference_kernel[grid]( + q, + k, + v, + s, + g, + o, + hk0, + hv0, + hkt, + hvt, + scale=scale, + K=K, + V=V, + M=M, + BK=BK, + BV=BV, + NG=NG, + ) + return o, (hkt, hvt) + + +def fused_recurrent_gsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + scale: float = 1., + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, tuple[torch.Tensor]]: + B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + HQ = q.shape[2] + if HQ != H: + raise ValueError("GQA not supported yet.") + + BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) + NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) + + hk0, hv0 = None, None + if initial_state != (None, None) and initial_state is not None: + hk0, hv0 = initial_state + hkt, hvt = None, None + if output_final_state: + hkt, hvt = q.new_empty(N, H, K, M, dtype=torch.float), q.new_empty(N, H, M, V, dtype=torch.float) + + ok = q.new_empty(NK, *s.shape, dtype=torch.float) + gk, gv = None, g + grid = (NM, NK, N * H) + fused_recurrent_fwd_kernel[grid]( + q=q, + k=k, + v=s, + g=None, + g_gamma=None, + gk=gk, + gv=gv, + o=ok, + h0=hk0, + ht=hkt, + cu_seqlens=cu_seqlens, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=M, + BK=BK, + BV=BM, + USE_G=False, + USE_G_GAMMA=False, + USE_GK=False, + USE_GV=True, + REVERSE=reverse, + ) + ok = ok.sum(0) + + qv = ok.softmax(-1, dtype=torch.float) + ov = q.new_empty(NM, *v.shape, dtype=torch.float) + gk, gv = g, None + grid = (NV, NM, N * H) + fused_recurrent_fwd_kernel[grid]( + q=qv, + k=s, + v=v, + g=None, + g_gamma=None, + gk=gk, + gv=gv, + o=ov, + h0=hv0, + ht=hvt, + cu_seqlens=cu_seqlens, + scale=1., + B=B, + T=T, + H=H, + K=M, + V=V, + BK=BM, + BV=BV, + USE_G=False, + USE_G_GAMMA=False, + USE_GK=True, + USE_GV=False, + REVERSE=reverse, + ) + ov = ov.sum(0) + return ok, hkt, qv, ov, hvt + + +def fused_recurrent_gsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + qv: torch.Tensor, + hk0: torch.Tensor | None = None, + hv0: torch.Tensor | None = None, + ok: torch.Tensor | None = None, + do: torch.Tensor | None = None, + dhkt: torch.Tensor | None = None, + dhvt: torch.Tensor | None = None, + scale: float = 1., + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor]: + B, T, H, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) + NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) + + dqv = q.new_empty(NV, B, T, H, M, dtype=torch.float) + dsv = q.new_empty(NV, B, T, H, M, dtype=torch.float) + dv = q.new_empty(NM, B, T, H, V, dtype=torch.float) + dgv = q.new_empty(NV, B, T, H, M, dtype=torch.float) + dhv0 = torch.empty_like(hv0)if hv0 is not None else None + + grid = (NV, NM, N * H) + fused_recurrent_bwd_kernel[grid]( + q=qv, + k=s, + v=v, + g=None, + g_gamma=None, + gk=g, + gv=None, + o=None, + h0=hv0, + do=do, + dq=dqv, + dk=dsv, + dv=dv, + dg=None, + dgk=dgv, + dgv=None, + dht=dhvt, + dh0=dhv0, + cu_seqlens=cu_seqlens, + scale=1., + B=B, + T=T, + H=H, + K=M, + V=V, + BK=BM, + BV=BV, + USE_G=False, + USE_G_GAMMA=False, + USE_GK=True, + USE_GV=False, + REVERSE=reverse, + ) + dqv = dqv.sum(0) + dsv = dsv.sum(0) + dv = dv.sum(0) + dgv = dgv.sum(0) + + dok = qv * (dqv - (qv * dqv).sum(-1, True)) + dq = q.new_empty(NM, B, T, H, K, dtype=torch.float) + dk = q.new_empty(NM, B, T, H, K, dtype=torch.float) + dsk = q.new_empty(NK, B, T, H, M, dtype=torch.float) + dgk = q.new_empty(NK, B, T, H, M, dtype=torch.float) + dhk0 = torch.empty_like(hk0)if hk0 is not None else None + + grid = (NM, NK, N * H) + fused_recurrent_bwd_kernel[grid]( + q=q, + k=k, + v=s, + g=None, + g_gamma=None, + gk=None, + gv=g, + o=ok, + h0=hk0, + do=dok, + dq=dq, + dk=dk, + dv=dsk, + dg=None, + dgk=None, + dgv=dgk, + dht=dhkt, + dh0=dhk0, + cu_seqlens=cu_seqlens, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=M, + BK=BK, + BV=BM, + USE_G=False, + USE_G_GAMMA=False, + USE_GK=False, + USE_GV=True, + REVERSE=reverse, + ) + dq = dq.sum(0) + dk = dk.sum(0) + dsk = dsk.sum(0) + dgk = dgk.sum(0) + + ds = dsk.add_(dsv) + dg = dgk.add_(dgv) + + return dq, dk, dv, ds, dg, dhk0, dhv0 + + +class FusedRecurrentGSAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, + hk0: torch.Tensor | None = None, + hv0: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor]]: + T = q.shape[1] + if T == 1 and not q.requires_grad: + o, (hkt, hvt) = fused_recurrent_gsa_inference( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + ) + return o, hkt, hvt + ok, hkt, qv, ov, hvt = fused_recurrent_gsa_fwd( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, k, v, s, g, qv, hk0, hv0, ok) + ctx.scale = scale + ctx.reverse = reverse + ctx.cu_seqlens = cu_seqlens + return ov.to(q.dtype), hkt, hvt + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dhkt=None, dhvt=None): + q, k, v, s, g, qv, hk0, hv0, ok = ctx.saved_tensors + scale = ctx.scale + reverse = ctx.reverse + cu_seqlens = ctx.cu_seqlens + + dq, dk, dv, ds, dg, dhk0, dhv0 = fused_recurrent_gsa_bwd( + q=q, + k=k, + v=v, + s=s, + g=g, + qv=qv, + hk0=hk0, + hv0=hv0, + ok=ok, + do=do, + dhkt=dhkt, + dhvt=dhvt, + scale=scale, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + return dq.to(q), dk.to(k), dv.to(v), ds.to(s), dg.to(g), None, dhk0, dhv0, None, None, None + + +def fused_recurrent_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor | None = None, + scale: int | None = None, + initial_state: tuple[torch.Tensor] | None = None, + output_final_state: bool | None = False, + reverse: bool | None = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + s (torch.Tensor): + slot representations of shape `[B, T, H, M]`. + g (torch.Tensor): + Forget gates of shape `[B, H, T, M]` applied to keys. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[Tuple[torch.Tensor]]): + Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]` and `[N, H, M, V]`. + Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (Tuple[torch.Tensor]): + Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gsa import fused_recurrent_gsa + # inputs with equal lengths + >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> s = torch.randn(B, T, H, M, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) + >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) + >>> o, (hk, hv) = fused_recurrent_gsa( + q, k, v, s, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, (hk_var, hv_var) = fused_recurrent_gsa( + q, k, v, s, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert hk.allclose(hk_var) + >>> assert hv.allclose(hv_var) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if initial_state is None: + initial_state = (None, None) + o, *final_state = FusedRecurrentGSAFunction.apply( + q, + k, + v, + s, + g, + scale, + *initial_state, + output_final_state, + reverse, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/gsa/naive.py b/fla/ops/gsa/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..f306aada68d15a416a04abf1cbe2785dd482a47a --- /dev/null +++ b/fla/ops/gsa/naive.py @@ -0,0 +1,67 @@ + + +import torch +from einops import repeat + + +def naive_recurrent_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor | None = None, + scale: int | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, +) -> torch.Tensor: + dtype = q.dtype + q, k, v, s, g = map(lambda x: x.transpose(1, 2).contiguous().float(), (q, k, v, s, g)) + + NG = q.shape[1]//k.shape[1] + # [batch_size, n_heads, seq_len, n_slots] + if g is None: + z = s.float().logcumsumexp(2) + g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z + s = torch.exp(s - z) + k, v, s, g = map(lambda x: repeat(x, 'b h t d -> b (h g) t d', g=NG), (k, v, s, g)) + if initial_state is not None: + initial_state = tuple(map(lambda x: repeat(x, 'b h k v -> b (h g) k v', g=NG), initial_state)) + + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + + hk = torch.zeros(B, H, K, M, dtype=torch.float, device=q.device) + ok = torch.zeros_like(s) + + if scale is None: + scale = q.shape[-1] ** -0.5 + + final_state = None + if initial_state is not None: + hk += initial_state[0] + + for i in range(T): + q_i = q[:, :, i] * scale + k_i = k[:, :, i] + v_i = s[:, :, i] + g_i = g[:, :, i].exp() + hk = hk * g_i[..., None, :] + k_i[..., None] * v_i[..., None, :] + ok[:, :, i] = (q_i[..., None] * hk).sum(-2) + + qv = ok.softmax(-1) + hv = torch.zeros(B, H, M, V, dtype=torch.float, device=q.device) + ov = torch.zeros_like(v) + if initial_state is not None: + hv += initial_state[1] + + for i in range(T): + q_i = qv[:, :, i] + k_i = s[:, :, i] + v_i = v[:, :, i] + g_i = g[:, :, i].exp() + hv = hv * g_i[..., :, None] + k_i[..., None] * v_i[..., None, :] + ov[:, :, i] = (q_i[..., None] * hv).sum(-2) + + if output_final_state: + final_state = (hk.view(B, -1, NG, K, M)[:, :, 0], hv.view(B, -1, NG, M, V)[:, :, 0]) + ov = ov.transpose(1, 2).contiguous() + return ov.to(dtype), final_state diff --git a/fla/ops/hgrn/__init__.py b/fla/ops/hgrn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..575f5902029609449f105d400352bec9487e1428 --- /dev/null +++ b/fla/ops/hgrn/__init__.py @@ -0,0 +1,8 @@ + +from .chunk import chunk_hgrn +from .fused_recurrent import fused_recurrent_hgrn + +__all__ = [ + 'chunk_hgrn', + 'fused_recurrent_hgrn', +] diff --git a/fla/ops/hgrn/chunk.py b/fla/ops/hgrn/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..356275d6702eaa72f81e56e50a36a17ab4506cea --- /dev/null +++ b/fla/ops/hgrn/chunk.py @@ -0,0 +1,282 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# this function implements the chunkwise form of HGRN, inspired by +# [Volodymyr Kyrylov in his blog post](https://proger.github.io/posts/scan/chunk.html) +# also refer to the `accelerated-scan` lib: https://github.com/proger/accelerated-scan + +# from tests on H800, with B, D = 16, 128, we see that the chunk can be greatly faster than the recurrent: +# +# Performance: +# seq_len chunk recurrent chunk_bwd recurrent_bwd +# 0 128.0 0.039360 0.061056 0.312160 0.205008 +# 1 256.0 0.045824 0.123712 0.308784 0.297696 +# 2 512.0 0.058688 0.241952 0.310720 0.626528 +# 3 1024.0 0.088288 0.476992 0.313184 1.333152 +# 4 2048.0 0.169472 0.943264 0.452464 2.724864 +# 5 4096.0 0.329920 1.886144 0.881600 5.551520 +# 6 8192.0 0.647872 3.755040 1.740496 11.117184 +# 7 16384.0 1.272064 7.520576 3.446608 22.362528 + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, input_guard + + +@triton.autotune( + configs=[ + triton.Config({'BD': 32}, num_warps=1), + triton.Config({'BD': 32}, num_warps=2), + triton.Config({'BD': 32}, num_warps=4), + triton.Config({'BD': 32}, num_warps=8), + triton.Config({'BD': 64}, num_warps=1), + triton.Config({'BD': 64}, num_warps=2), + triton.Config({'BD': 64}, num_warps=4), + triton.Config({'BD': 64}, num_warps=8), + triton.Config({'BD': 128}, num_warps=1), + triton.Config({'BD': 128}, num_warps=2), + triton.Config({'BD': 128}, num_warps=4), + triton.Config({'BD': 128}, num_warps=8), + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_hgrn_fwd_kernel_h( + x, + g, + gc, + o, + h0, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, +): + i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + p_x = x + i_b * T * D + i_t * BT * D + o_d + p_g = g + i_b * T * D + i_t * BT * D + o_d + p_gc = gc + i_b * T * D + i_t * BT * D + o_d + p_o = o + i_b * T * D + i_t * BT * D + o_d + + b_h = tl.zeros([BD], dtype=tl.float32) + b_gc = tl.zeros([BD], dtype=tl.float32) + if USE_INITIAL_STATE: + if i_t == 0: + b_h += tl.load(h0 + i_b * D + o_d, mask=mask, other=0).to(tl.float32) + for i in range(0, BT): + mask_t = mask & ((i_t * BT + i) < T) + b_x = tl.load(p_x, mask=mask_t, other=0).to(tl.float32) + b_g = tl.load(p_g, mask=mask_t, other=0).to(tl.float32) + b_h = exp(b_g) * b_h + b_x + b_gc = b_gc + b_g + tl.store(p_gc, b_gc.to(p_o.dtype.element_ty), mask=mask_t) + tl.store(p_o, b_h.to(p_o.dtype.element_ty), mask=mask_t) + + p_x += D + p_g += D + p_gc += D + p_o += D + + +@triton.jit(do_not_specialize=['T']) +def chunk_hgrn_fwd_kernel_o( + gc, + o, + s_b, + s_t, + s_d, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_d, i_b = tl.program_id(0), tl.program_id(1) + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + for i_t in range(1, tl.cdiv(T, BT)): + p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + + # [BD,] + b_h0 = tl.load(o + i_b * T * D + i_t * BT * D - D + o_d, mask=mask, other=0).to(tl.float32) + # [BT, BD] + b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_o = b_o + exp(b_gc) * b_h0[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [32, 64, 128] + for num_warps in [1, 2, 4, 8] + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_hgrn_bwd_kernel_h( + g, + gc, + dx, + do, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + BC = min(BT, T - i_t * BT) + NT = tl.num_programs(1) + + p_g = g + (i_b * T + i_t * BT + BC - 1) * D + o_d + p_gc = gc + (i_b * T + i_t * BT + BC - 1) * D + o_d + p_dx = dx + (i_b * T + i_t * BT + BC - 1) * D + o_d + p_do = do + (i_b * T + i_t * BT + BC - 1) * D + o_d + + if i_t == NT - 1: + b_gc = tl.zeros([BD], dtype=tl.float32) + else: + b_gc = tl.load(g + (i_b * T + i_t * BT + BT) * D + o_d, mask=mask, other=0).to(tl.float32) + b_dh = tl.zeros([BD], dtype=tl.float32) + for _ in range(BC - 1, -1, -1): + tl.store(p_gc, b_gc.to(p_gc.dtype.element_ty), mask=mask) + + b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask, other=0).to(tl.float32) + + b_gc = b_gc + b_g + b_dh = b_dh + b_do + b_dx = b_dh + b_dh = b_dh * exp(b_g) + + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=mask) + + p_g -= D + p_gc -= D + p_dx -= D + p_do -= D + + +@triton.jit(do_not_specialize=['T']) +def chunk_hgrn_bwd_kernel_o( + g, + gc, + o, + dx, + dg, + s_b, + s_t, + s_d, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_d, i_b = tl.program_id(0), tl.program_id(1) + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): + p_g = tl.make_block_ptr(g + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT - 1, i_d * BD), (BT, BD), (1, 0)) + p_dx = tl.make_block_ptr(dx + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_dg = tl.make_block_ptr(dg + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + + # [BD,] + mask_t = mask & ((i_t + 1) * BT < T) + b_ht = tl.load(dx + i_b * T * D + (i_t + 1) * BT * D + o_d, mask=mask_t, other=0).to(tl.float32) + # [BT, BD] + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_dx = tl.load(p_dx, boundary_check=(0, 1)).to(tl.float32) + + b_dx = b_dx + exp(b_gc) * b_ht[None, :] + b_dg = b_o * b_dx * exp(b_g) + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +class ChunkHGRNFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward(ctx, x, g, initial_state=None, output_final_state=False): + B, T, D = x.shape + BT, BD = 128, min(64, triton.next_power_of_2(D)) + num_warps = 8 if BD == 64 else 4 + + gc = torch.empty_like(g, dtype=torch.float) + o = torch.empty_like(x, dtype=torch.float) + def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) + chunk_hgrn_fwd_kernel_h[grid]( + x, g, gc, o, initial_state, + T=T, D=D, BT=BT, + USE_INITIAL_STATE=initial_state is not None, + ) + def grid(meta): return (triton.cdiv(D, meta['BD']), B) + chunk_hgrn_fwd_kernel_o[grid]( + gc, o, + o.stride(-3), o.stride(-2), o.stride(-1), + T=T, D=D, BT=BT, BD=BD, + num_warps=num_warps, + ) + final_state = None + if output_final_state: + final_state = o[:, -1].clone() + o = o.to(x.dtype) + ctx.save_for_backward(g, o, initial_state) + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht=None): + g, o, initial_state = ctx.saved_tensors + B, T, D = do.shape + BT, BD = 128, min(64, triton.next_power_of_2(D)) + num_warps = 8 if BD == 64 else 4 + + gc = torch.empty_like(g, dtype=torch.float) + dx = torch.empty_like(o, dtype=torch.float) + def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) + chunk_hgrn_bwd_kernel_h[grid]( + g, gc, dx, do, + T=T, D=D, BT=BT, + ) + + dg = torch.empty_like(g, dtype=torch.float) + def grid(meta): return (triton.cdiv(D, meta['BD']), B) + chunk_hgrn_bwd_kernel_o[grid]( + g, gc, o, dx, dg, + o.stride(-3), o.stride(-2), o.stride(-1), + T=T, D=D, BT=BT, BD=BD, + num_warps=num_warps, + ) + if initial_state is not None: + dg[:, 0] = (initial_state * dx[:, 0] * g[:, 0].float().exp()).to(dg.dtype) + + return dx.to(o.dtype), dg, None, None + + +@torch.compiler.disable +def chunk_hgrn( + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + return ChunkHGRNFunction.apply(x, g, initial_state, output_final_state) diff --git a/fla/ops/hgrn/fused_recurrent.py b/fla/ops/hgrn/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..21680c73237b2218379095c07492bd59b034f5ae --- /dev/null +++ b/fla/ops/hgrn/fused_recurrent.py @@ -0,0 +1,308 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [32, 64, 128] + for num_warps in [1, 2, 4, 8] + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_hgrn_fwd_kernel( + x, + g, + o, + h0, + ht, + cu_seqlens, + T, + D: tl.constexpr, + BD: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_n = tl.program_id(0), tl.program_id(1) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + p_x = x + bos * D + o_d + p_g = g + bos * D + o_d + p_o = o + bos * D + o_d + + b_h = tl.zeros([BD], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_n * D + o_d + b_h += tl.load(p_h0, mask=mask, other=0).to(tl.float32) + for _ in range(0, T): + b_x = tl.load(p_x, mask=mask, other=0).to(tl.float32) + b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) + b_h = exp(b_g) * b_h + b_x + tl.store(p_o, b_h.to(p_o.dtype.element_ty), mask=mask) + + p_x += D + p_g += D + p_o += D + + if STORE_FINAL_STATE: + p_ht = ht + i_n * D + o_d + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [32, 64, 128] + for num_warps in [1, 2, 4, 8] + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_hgrn_bwd_kernel( + g, + o, + h0, + dx, + dg, + do, + dht, + dh0, + cu_seqlens, + T, + D: tl.constexpr, + BD: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_n = tl.program_id(0), tl.program_id(1) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + p_g = g + (bos + T - 1) * D + o_d + p_o = o + (bos + T - 2) * D + o_d + p_dx = dx + (bos + T - 1) * D + o_d + p_dg = dg + (bos + T - 1) * D + o_d + p_do = do + (bos + T - 1) * D + o_d + + b_dh = tl.zeros([BD], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = dht + i_n * D + o_d + b_dh += tl.load(p_dht, mask=mask, other=0).to(tl.float32) + + for i in range(T - 1, -1, -1): + b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask, other=0).to(tl.float32) + if i > 0: + b_o = tl.load(p_o, mask=mask, other=0).to(tl.float32) + elif USE_INITIAL_STATE: + b_o = tl.load(h0 + i_n * D + o_d, mask=mask, other=0).to(tl.float32) + else: + b_o = tl.zeros([BD], dtype=tl.float32) + + b_dh = b_dh + b_do + b_dx = b_dh + b_dh = b_dh * exp(b_g) + b_dg = b_dh * b_o + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=mask) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=mask) + + p_g -= D + p_o -= D + p_dx -= D + p_dg -= D + p_do -= D + + if USE_INITIAL_STATE: + p_dh0 = dh0 + i_n * D + o_d + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask) + + +def fused_recurrent_hgrn_fwd( + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, D = x.shape + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + o = torch.empty_like(x) + final_state = x.new_empty(N, D) if output_final_state else None + + def grid(meta): return (triton.cdiv(D, meta['BD']), N) + fused_recurrent_hgrn_fwd_kernel[grid]( + x=x, + g=g, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + T=T, + D=D, + ) + return o, final_state + + +def fused_recurrent_hgrn_bwd( + g: torch.Tensor, + o: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor = None, + initial_state: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, D = do.shape + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + dx = torch.empty_like(o, dtype=torch.float) + dg = torch.empty_like(g, dtype=torch.float) + dh0 = torch.empty_like(initial_state, dtype=torch.float) if initial_state is not None else None + def grid(meta): return (triton.cdiv(D, meta['BD']), N) + fused_recurrent_hgrn_bwd_kernel[grid]( + g=g, + o=o, + h0=initial_state, + dx=dx, + dg=dg, + do=do, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + T=T, + D=D, + ) + return dx, dg, dh0 + + +class FusedRecurrentHGRNFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, ht = fused_recurrent_hgrn_fwd( + x=x, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(g, o, initial_state) + ctx.cu_seqlens = cu_seqlens + return o, ht + + @staticmethod + @input_guard + def backward(ctx, do, dht=None): + g, o, initial_state = ctx.saved_tensors + cu_seqlens = ctx.cu_seqlens + + dx, dg, dh0 = fused_recurrent_hgrn_bwd( + g=g, + o=o, + do=do, + dht=dht, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ) + return dx, dg, dh0, None, None + + +@torch.compiler.disable +def fused_recurrent_hgrn( + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + x (torch.Tensor): + inputs of shape `[B, T, D]. + g (torch.Tensor): + Forget gates of shape `[B, T, D]`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, D]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, D]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, D]`. + final_state (torch.Tensor): + Final state of shape `[N, D]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.hgrn import fused_recurrent_hgrn + # inputs with equal lengths + >>> B, T, D = 4, 2048, 512 + >>> x = torch.randn(B, T, D, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, D, device='cuda')) + >>> h0 = torch.randn(B, D, device='cuda') + >>> o, ht = fused_recurrent_hgrn(x, g, initial_state=h0, output_final_state=True) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> x, g = map(lambda x: rearrange(x, 'b t d -> 1 (b t) d'), (x, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = x.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_hgrn(x, g, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert ht.allclose(ht_var) + """ + return FusedRecurrentHGRNFunction.apply( + x, + g, + initial_state, + output_final_state, + cu_seqlens, + ) diff --git a/fla/ops/hgrn/naive.py b/fla/ops/hgrn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..ce2cf030edb9b9c11ed686f33f4e5e7ccbf8d409 --- /dev/null +++ b/fla/ops/hgrn/naive.py @@ -0,0 +1,61 @@ + + +import torch + + +def naive_recurrent_hgrn( + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, +) -> torch.Tensor: + dtype = x.dtype + x, g = map(lambda i: i.float(), (x, g)) + B, T, D = x.shape + + h = torch.zeros(B, D, dtype=torch.float, device=x.device) + o = torch.zeros_like(x) + + final_state = None + if initial_state is not None: + h += initial_state + + for i in range(T): + h = g[:, i].exp() * h + x[:, i] + o[:, i] = h + + if output_final_state: + final_state = h + return o.to(dtype), final_state + + +def naive_chunk_hgrn( + x: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, + chunk_size: int = 64, +) -> torch.Tensor: + dtype = x.dtype + x, g = map(lambda i: i.float(), (x, g)) + B, T, D = x.shape + + gc = g.view(B, chunk_size, D).cumsum(-2).view_as(g) + h = torch.zeros(B, D, dtype=torch.float, device=x.device) + o = torch.zeros_like(x) + + final_state = None + if initial_state is not None: + h += initial_state + + for i in range(0, T, chunk_size): + hp = h + h = torch.zeros(B, D, dtype=torch.float, device=x.device) + for j in range(i, i + chunk_size): + h = g[:, j].exp() * h + x[:, j] + o[:, j] = hp * gc[:, j].exp() + h + h = o[:, j].clone() + + if output_final_state: + final_state = h + return o.to(dtype), final_state diff --git a/fla/ops/kda/__init__.py b/fla/ops/kda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..483321f63b4702705ac4f5981faffb06f5c86033 --- /dev/null +++ b/fla/ops/kda/__init__.py @@ -0,0 +1,7 @@ +from .chunk import chunk_kda +from .fused_recurrent import fused_recurrent_kda + +__all__ = [ + "chunk_kda", + "fused_recurrent_kda", +] diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..b21b8a4a6e46a1d6777fb115b86d5d941e52be69 --- /dev/null +++ b/fla/ops/kda/chunk.py @@ -0,0 +1,347 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang +# Related files are modified and supported by the Moonshot AI Team + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.cp import FLACPContext +from fla.ops.kda.chunk_bwd import chunk_kda_bwd +from fla.ops.kda.chunk_fwd import chunk_kda_fwd +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +class ChunkKDAFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + ): + chunk_size = 64 + + # Apply l2norm + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g_input = g + + (o, final_state, g_cumsum, Aqk, Akk, w, u, qg, kg, v_new, h, initial_state) = chunk_kda_fwd( + q=q, + k=k, + v=v, + g=g_input, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + disable_recompute=disable_recompute, + return_intermediate_states=return_intermediate_states, + cp_context=cp_context, + transpose_state_layout=transpose_state_layout, + ) + + if return_intermediate_states: + assert torch.is_inference_mode_enabled(), "return_intermediate_states is only allowed in inference mode" + assert disable_recompute is False, "return_intermediate_states must be used with disable_recompute=False" + return o.type_as(q), final_state, h + + ctx.save_for_backward( + q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta, A_log, dt_bias, Aqk, Akk, + w, u, qg, kg, v_new, h, + initial_state, cu_seqlens, chunk_indices + ) + ctx.chunk_size = chunk_size + ctx.safe_gate = safe_gate + ctx.scale = scale + ctx.lower_bound = lower_bound + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.disable_recompute = disable_recompute + ctx.cp_context = cp_context + ctx.transpose_state_layout = transpose_state_layout + return o.type_as(q), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + (q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta, A_log, dt_bias, Aqk, Akk, + w, u, qg, kg, v_new, h, + initial_state, cu_seqlens, chunk_indices) = ( + ctx.saved_tensors + ) + + dq, dk, dv, db, dg, dh0, dA, dbias = chunk_kda_bwd( + q=q, + k=k, + v=v, + g=g_cumsum, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=ctx.chunk_size, + safe_gate=ctx.safe_gate, + g_org=g_input if ctx.use_gate_in_kernel else None, lower_bound=ctx.lower_bound, + use_gate_in_kernel=ctx.use_gate_in_kernel, + A_log=A_log, dt_bias=dt_bias, + disable_recompute=ctx.disable_recompute, + w=w, u=u, qg=qg, kg=kg, v_new=v_new, h=h, + cp_context=ctx.cp_context, + transpose_state_layout=ctx.transpose_state_layout, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + + return (dq.to(q), dk.to(k), dv.to(v), dg.to(g_input), db.to(beta), dA, dbias, None, dh0, + None, None, None, None, None, None, None, None, None, None, None) + + +@torch.compiler.disable +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context: FLACPContext = None, + transpose_state_layout: bool = False, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H, K]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the KDA attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q,k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space KDA decay internally. + - If `True`: + The passed `g` acts as the raw input for `-exp(A_log).view(H, -1) * softplus(g + dt_bias.view(H, K))`. + Note that as part of the input arguments, + `A_log` (shape `[H]`) and the optional `dt_bias` (shape `[H * K]`) should be provided. + - If `False`, `g` is expected to be the pre-computed decay value. + Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cu_seqlens_cpu (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + safe_gate (bool): + Whether the kernel can assume the input gate values `g` are in a safe range. + When `True`, the kernel can use M=16 TensorCore acceleration. + The safe range is approximately [-5, 0). Default: `False`. + lower_bound (Optional[float]): + Lower bound for the forget gate activation function when `use_gate_in_kernel=True`. + This parameter modifies the internal forget gate activation and is recommended + to be set to `-5` when `safe_gate` is enabled. Default: `None`. + disable_recompute (bool): + Whether to disable gradient recomputation in the kernel. When `True`, the kernel + will save all intermediate activations for backward pass, which is beneficial + for training small models at the cost of increased memory usage. Default: `False`. + return_intermediate_states (bool): + If True, returns intermediate state `h` for inference scenarios (e.g., vLLM). + Must be used within `torch.inference_mode()` and will return a 3-tuple instead of 2-tuple. + This is not intended for training as it bypasses autograd. Default: `False`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + transpose_state_layout (Optional[bool]): + Whether to use the transposed state layout for the hidden state. + Default: `False`. + + Returns: + - Normal mode (return_intermediate_states=False): A tuple (o, final_state) + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + - Inference mode (return_intermediate_states=True): A tuple (o, final_state, h) + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + h (torch.Tensor): + Intermediate states of shape `[B, NT, H, K, V]` and dtype `bfloat16` for caching or further processing. + - For equal-length sequences: `NT = #chunks_per_sequence` (typically `ceil(T / chunk_size)`) + - For variable-length sequences (cu_seqlens): B is always 1 (flattened), + NT is the total number of chunks across all sequences, + determined by `prepare_chunk_indices(cu_seqlens, chunk_size)` + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.kda import chunk_kda + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda') + >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda') + >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda') + >>> o, ht = chunk_kda( + q, k, v, g, beta, + A_log=A_log, + dt_bias=dt_bias, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_kda( + q, k, v, g, beta, + A_log=A_log, + dt_bias=dt_bias, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + # Override cu_seqlens and cu_seqlens_cpu with the ones from the context + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if initial_state is not None: + assert initial_state.dtype == torch.float32, "initial_state must be in float32." + + A_log, dt_bias = None, None + if use_gate_in_kernel: + assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." + A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + + if safe_gate and use_gate_in_kernel: + if lower_bound is None: + raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") + + assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." + assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" + assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." + assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + + if scale is None: + scale = k.shape[-1] ** -0.5 + return ChunkKDAFunction.apply( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + cu_seqlens, + cu_seqlens_cpu, + safe_gate, + lower_bound, + disable_recompute, + return_intermediate_states, + cp_context, + transpose_state_layout, + ) diff --git a/fla/ops/kda/chunk_bwd.py b/fla/ops/kda/chunk_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..d9d528b6fafe3aead37a857b0f2deb703fcc035f --- /dev/null +++ b/fla/ops/kda/chunk_bwd.py @@ -0,0 +1,583 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + expand_h0, +) +from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra +from fla.ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum +from fla.ops.kda.wy_fast import recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_bwd_kernel_dAv( + q, + k, + v, + A, + do, + dv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + dA += (bos * H + i_h) * BT + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BV, BT] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT, BT] + b_dA += tl.dot(b_do, b_v) + # [BT, BV] + b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['BT', 'TRANSPOSE_STATE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_bwd_kernel_wy_dqkg_fused( + q, + k, + v, + v_new, + g, + beta, + A, + h, + do, + dh, + dq, + dk, + dv, + dv2, + dg, + db, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t.to(tl.int64) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_last = (o_t == min(T, i_t * BT + BT) - 1) + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + g += (bos * H + i_h) * K + beta += bos * H + i_h + A += (bos * H + i_h) * BT + h += (i_tg * H + i_h) * K*V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K*V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dv2 += (bos * H + i_h) * V + dg += (bos * H + i_h) * K + db += bos * H + i_h + dA += (bos * H + i_h) * BT + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_db = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + + p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v_new = tl.make_block_ptr(v_new, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + if TRANSPOSE_STATE: + p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + else: + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BT, BV] + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + + b_dgk += tl.sum(b_h * b_dh, axis=0) + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) + b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) + tl.debug_barrier() # DO NOT REMOVE THIS LINE! + if i_k == 0: + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_v = tl.load(p_v, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dv, tl.trans(b_v)) + + b_dvb = tl.dot(b_A, b_dv) + b_dv2 = b_dvb * b_beta[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + + b_gk_exp = exp2(b_g) + b_gb = b_gk_exp * b_beta[:, None] + b_dgk *= exp2(b_gn) + b_dq = b_dq * b_gk_exp * scale + b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0) + + b_kg = b_k * b_gk_exp + + b_dw = -b_dw.to(b_A.dtype) + b_dA += tl.dot(b_dw, tl.trans(b_kg.to(b_A.dtype))) + + b_dkgb = tl.dot(b_A, b_dw) + b_db += tl.sum(b_dkgb * b_kg, 1) + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kdk = b_k * b_dk + b_dgk += tl.sum(b_kdk, axis=0) + b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_beta[:, None] + b_dk = b_dk + b_dkgb * b_gb + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + b_dA = tl.where(m_A, -b_dA, 0) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def chunk_kda_bwd_dAv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + A: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # H100 can have larger block size + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem: + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dA = v.new_empty(B, T, H, BT, dtype=torch.float) + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_kda_bwd_kernel_dAv[grid]( + q=q, + k=k, + v=v, + A=A, + do=do, + dv=dv, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dA, dv + + +def chunk_kda_bwd_wy_dqkg_fused( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + transpose_state_layout: bool = False, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dv2 = torch.empty_like(v) + dg = torch.empty_like(g, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + + grid = (NT, B * H) + chunk_kda_bwd_kernel_wy_dqkg_fused[grid]( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dv=dv, + dv2=dv2, + dg=dg, + db=db, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + TRANSPOSE_STATE=transpose_state_layout, + ) + dv = dv2 + return dq, dk, dv, db, dg, dA + + +def chunk_kda_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + g: torch.Tensor | None = None, + g_org: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + **kwargs, +): + if disable_recompute is False: + if use_gate_in_kernel: + g = kda_gate_chunk_cumsum( + g=g_org, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound + ) + w, u, qg, kg = recompute_w_u_fwd( + q=q, + k=k, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if cp_context is not None: + # Restore the full initial_state tensor from the compressed version. + # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. + initial_state = expand_h0(initial_state, context=cp_context) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + else: + w, u, qg, kg, v_new, h = kwargs["w"], kwargs["u"], kwargs["qg"], kwargs["kg"], kwargs["v_new"], kwargs["h"] + if cp_context is not None: + # Restore the full initial_state tensor from the compressed version. + # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. + initial_state = expand_h0(initial_state, context=cp_context) + + # dAqk = do @ v.T + # dv = A @ do + dAqk, dv = chunk_kda_bwd_dAv( + q=q, + k=k, + v=v_new, + do=do, + A=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=qg, + k=kg, + w=w, + do=do, + dv=dv, + gk=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + use_exp2=True, + context=cp_context, + transpose_state_layout=transpose_state_layout, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, + k=kg, + w=w, + gk=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + + dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused( + q=q, + k=k, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=Akk, + h=h, + do=do, + dh=dh, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + transpose_state_layout=transpose_state_layout, + ) + + dq, dk, db, dg = chunk_kda_bwd_intra( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dk=dk, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=safe_gate + ) + + dA, dbias = None, None + dg = chunk_local_cumsum( + dg, + chunk_size=chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + if use_gate_in_kernel: + dg, dA, dbias = kda_gate_bwd( + g=g_org, + A_log=A_log, + dt_bias=dt_bias, + dyg=dg, + lower_bound=lower_bound + ) + + return dq, dk, dv, db, dg, dh0, dA, dbias diff --git a/fla/ops/kda/chunk_fwd.py b/fla/ops/kda/chunk_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..fa01a6fd28e286bde8212795c2d52b098404792c --- /dev/null +++ b/fla/ops/kda/chunk_fwd.py @@ -0,0 +1,132 @@ +import torch + +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, +) +from fla.ops.gla.chunk import chunk_gla_fwd_o_gk +from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, +): + # Apply gate activation + g_org = None + if use_gate_in_kernel: + g_org = g + g = kda_gate_chunk_cumsum( + g=g_org, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + else: + g = chunk_local_cumsum( + g=g, + scale=RCP_LN2, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices + ) + + # qg = None if disable_recompute is False + w, u, qg, kg, Aqk, Akk = chunk_kda_fwd_intra( + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + disable_recompute=disable_recompute + ) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=kg, + w=w, + u=u, + gk=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + + if cp_context is not None: + # In Context Parallel (CP) mode, global initial states are not supported at the entry point. + # The `initial_state` here is computed internally via inter-rank communication. + # Since only the first sequence in the local batch can be a continuation of a cross-rank sequence, + # only the first state in the tensor is relevant. We compress it to optimize memory for `save_for_backward`. + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + use_exp2=True, + transpose_state_layout=transpose_state_layout, + ) + if disable_recompute is False: + # Delete to save memory + w, u, qg, kg, v_new = None, None, None, None, None + if not return_intermediate_states: + # Only delete h if not requested for inference + h = None + if use_gate_in_kernel: + g = None + return o, final_state, g, Aqk, Akk, w, u, qg, kg, v_new, h, initial_state diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd30016aaf75603732af95a4f30c395707414d6 --- /dev/null +++ b/fla/ops/kda/chunk_intra.py @@ -0,0 +1,900 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.kda.chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel +from fla.ops.kda.wy_fast import recompute_w_u_fwd +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2, gather +from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs + +if IS_TF32_SUPPORTED: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') +else: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "K", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + Aqk += (bos * H + i_h) * BT + Akk += (bos * H + i_h) * BT + Akkd += (bos * H + i_h) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b1 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b2 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b3 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_kda_bwd_kernel_intra( + q, + k, + g, + beta, + dAqk, + dAkk, + dq, + dq2, + dk, + dk2, + dg, + dg2, + db, + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, + SAFE_GATE: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_k, i_i = i_kc // NC, i_kc % NC + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + beta += bos * H + i_h + + dAqk += (bos * H + i_h) * BT + dAkk += (bos * H + i_h) * BT + dq += (bos * H + i_h) * K + dq2 += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dk2 += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + dg2 += (bos * H + i_h) * K + db += (i_k * all + bos) * H + i_h + + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) + b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) + if i_i > 0: + p_gn = g + i_ti * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp2(b_gn - b_gk) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + # [BC, BK] + b_dq2 += tl.dot(b_dAqk, b_kg) + b_dk2 += tl.dot(b_dAkk, b_kg) + b_gqn = exp2(b_g - b_gn) + b_dq2 *= b_gqn + b_dk2 *= b_gqn + + o_i = tl.arange(0, BC) + m_dA = (i_ti + o_i) < T + o_dA = (i_ti + o_i) * H*BT + i_i * BC + p_kj = k + i_ti * H*K + o_k + p_gkj = g + i_ti * H*K + o_k + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_qk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) + b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) + b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) + exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) + exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) + + b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) + b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) + # [BK] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_gqk = exp2(b_g - b_gkj[None, :]) + b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * b_gqk, 0.) + b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.) + + p_kj += H*K + p_gkj += H*K + + b_db = tl.sum(b_dk2 * b_k, 1) + b_dk2 *= b_b[:, None] + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_ti,), (BC,), (0,)) + + b_dg2 = b_q * b_dq2 + b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + tl.debug_barrier() + b_dkt = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = g + (min(i_ti + BC, T) - 1) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + for i_j in range(i_i + 1, NC): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + # [BC] + b_b = tl.load(p_b, boundary_check=(0,)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC, BK] + b_gkn = exp2(b_gk - b_gn) + b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) + b_kbg = b_kb * tl.where(m_j[:, None], b_gkn, 0) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dkt += tl.dot(b_dAqk, b_qg) + b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt *= exp2(b_gn - b_g) + o_dA = i_ti * H*BT + i_i * BC + o_i + p_qj = q + i_ti * H*K + o_k + p_kj = k + i_ti * H*K + o_k + p_gkj = g + i_ti * H*K + o_k + p_bj = beta + i_ti * H + + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_kk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) + b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) + # ensure numerical stability + b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) + exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) + exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) + + b_q_exp = b_q * exp_b_g_diag_kk + b_kb_exp = b_k * b_b[:, None] * exp_b_g_diag_kk + + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dAqk = tl.load(dAqk + o_dA + j * H*BT) + b_dAkk = tl.load(dAkk + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_gkq = exp2(b_gkj[None, :] - b_g) + b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.) + b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * b_gkq, 0.) + + p_qj += H*K + p_kj += H*K + p_gkj += H*K + p_bj += H + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg2 = tl.make_block_ptr(dg2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + + b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) + b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) + b_dk2 += b_dkt + + tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * H + i_h) * K + beta = beta + bos * H + i_h + Aqk = Aqk + (bos * H + i_h) * BT + Akk = Akk + (bos * H + i_h) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (H*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * H*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, + disable_recompute: bool = False, +): + B, T, H, K = k.shape + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) + + # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) + # Step 1: compute diagonal blocks into Akk_diag (fp32) + if safe_gate: + grid = (NT, NC, B * H) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + grid = (NT, B * H) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + ) + w, u, qg, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=Akk, + q=q if disable_recompute else None, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, qg, kg, Aqk, Akk + + +def chunk_kda_bwd_intra( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, +): + B, T, H, K = k.shape + BT = chunk_size + BC = min(16, BT) + BK = min(32, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq2 = torch.empty_like(q) + dk2 = torch.empty_like(k) + db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) + dg2 = torch.empty_like(dg, dtype=torch.float) + grid = (NK * NC, NT, B * H) + chunk_kda_bwd_kernel_intra[grid]( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dq2=dq2, + dk=dk, + dk2=dk2, + dg=dg, + dg2=dg2, + db=db2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + SAFE_GATE=safe_gate, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + dq = dq2 + dk = dk2 + db = db2.sum(0).add_(db) + dg = dg2 + + return dq, dk, db, dg diff --git a/fla/ops/kda/chunk_intra_token_parallel.py b/fla/ops/kda/chunk_intra_token_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..3640cce972534fbaf2c7366cbb5c0e71ffcbb4b4 --- /dev/null +++ b/fla/ops/kda/chunk_intra_token_parallel.py @@ -0,0 +1,169 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Token-parallel implementation of KDA intra chunk kernel + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp2 +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + q += bos * H*K + k += bos * H*K + g += bos * H*K + Aqk += bos * H*BT + Akk += bos * H*BC + beta += bos * H + + BK: tl.constexpr = triton.next_power_of_2(K) + o_h = tl.arange(0, BH) + o_k = tl.arange(0, BK) + m_h = (i_hg * BH + o_h) < H + m_k = o_k < K + + p_q = tl.make_block_ptr(q + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_g = tl.make_block_ptr(g + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_t * H, (H,), (1,), (i_hg * BH,), (BH,), (0,)) + # [BH, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_k = b_k * tl.load(p_beta, boundary_check=(0,)).to(tl.float32)[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + p_kj = tl.make_block_ptr(k + j * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_gj = tl.make_block_ptr(g + j * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + # [BH, BK] + b_kj = tl.load(p_kj, boundary_check=(0, 1)).to(tl.float32) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = b_kj * exp2(b_g - b_gj) + + b_kgj = tl.where(m_k[None, :], b_kgj, 0.0) + # [BH] + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * H*BT + (i_hg * BH + o_h) * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_h) + tl.store(Akk + i_t * H*BC + (i_hg * BH + o_h) * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_h) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, H, K] cumsum of gates + beta: [B, T, H] + Aqk: [B, T, H, BT] output tensor to write to + Akk: [B, T, H, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K = q.shape + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(H, meta['BH'])) + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/fla/ops/kda/fused_recurrent.py b/fla/ops/kda/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..c6f6f923a7380dfef23ee83312e869562dd13832 --- /dev/null +++ b/fla/ops/kda/fused_recurrent.py @@ -0,0 +1,434 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# This kernel is modified from the Decode kernel of the vllm gdn/kda model. + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import input_guard + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + o, + h0, + ht, + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_init_state_token: tl.constexpr, + stride_final_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + stride_indices_tok: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + num_stages: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + NK = tl.cdiv(K, BK) + i_k = pid % NK + pid_rest = pid // NK + + i_v = pid_rest % NV + i_nh = pid_rest // NV + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + if T == 0: + # no tokens to process for this sequence + return + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + + p_g = g + (bos * HV + i_hv) * K + o_k + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + if TRANSPOSE_STATE: + mask_h = mask_v[:, None] & mask_k[None, :] + else: + mask_h = mask_k[:, None] & mask_v[None, :] + + if TRANSPOSE_STATE: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if IS_CONTINUOUS_BATCHING: + if IS_SPEC_DECODING: + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 + p_h0 = ( + h0 + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_init_state_token + ) + if TRANSPOSE_STATE: + p_h0 = p_h0 + i_hv * K * V + o_v[:, None] * K + o_k[None, :] + else: + p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + else: + if TRANSPOSE_STATE: + p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] + else: + p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for i_t in tl.range(0, T, num_stages=num_stages): + b_q = tl.load(p_q, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0, eviction_policy='evict_first').to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + b_g = tl.load(p_g, eviction_policy='evict_last').to(tl.float32) + + if USE_GATE_IN_KERNEL: + b_A = tl.load(A_log + i_h).to(tl.float32) + + if HAS_DT_BIAS: + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_g = b_g + b_bias + + if USE_LOWER_BOUND: + b_gk = lower_bound * tl.sigmoid(exp(b_A) * b_g) + else: + b_gk = -exp(b_A) * softplus(b_g) + else: + b_gk = b_g + + if TRANSPOSE_STATE: + b_h *= exp(b_gk[None, :]) + else: + b_h *= exp(b_gk[:, None]) + + if TRANSPOSE_STATE: + b_v -= tl.sum(b_h * b_k[None, :], 1) + else: + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0, eviction_policy='evict_first').to(tl.float32) + else: + b_beta = tl.load(p_beta, eviction_policy='evict_last').to(tl.float32) + b_v *= b_beta + if TRANSPOSE_STATE: + b_h += b_v[:, None] * b_k[None, :] + b_o = tl.sum(b_h * b_q[None, :], 1) + else: + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v, eviction_policy='evict_first') + + if IS_CONTINUOUS_BATCHING: + if INPLACE_FINAL_STATE: + p_ht = ( + ht + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_final_state_token + ) + else: + p_ht = ht + (bos + i_t) * stride_final_state_token + if TRANSPOSE_STATE: + p_ht = p_ht + i_hv * K * V + o_v[:, None] * K + o_k[None, :] + else: + p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_g += HV * K + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + if not IS_CONTINUOUS_BATCHING: + if STORE_FINAL_STATE: + if TRANSPOSE_STATE: + p_ht = ht + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] + else: + p_ht = ht + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@torch.compiler.disable +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + scale: float | None = None, + output_final_state: bool = False, + inplace_final_state: bool = True, + cu_seqlens: torch.LongTensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + lower_bound: float | None = None, + out: torch.Tensor | None = None, + transpose_state_layout: bool = False, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + if scale is None: + scale = k.shape[-1] ** -0.5 + + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + BV = 32 + + if out is None: + out = torch.zeros_like(v) + else: + assert out.shape == v.shape + if inplace_final_state: + assert initial_state is not None + final_state = initial_state + elif output_final_state: + if transpose_state_layout: + final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) + else: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + stride_init_state_token = initial_state.stride(0) if initial_state is not None else 1 + stride_final_state_token = final_state.stride(0) if final_state is not None else 1 + + if ssm_state_indices is None: + stride_indices_seq, stride_indices_tok = 1, 1 + elif ssm_state_indices.ndim == 1: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 + else: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() + + grid = (triton.cdiv(V, BV) * N * HV, ) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + o=out, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + stride_init_state_token=stride_init_state_token, + stride_final_state_token=stride_final_state_token, + stride_indices_seq=stride_indices_seq, + stride_indices_tok=stride_indices_tok, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + INPLACE_FINAL_STATE=inplace_final_state, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + TRANSPOSE_STATE=transpose_state_layout, + num_warps=4, + num_stages=2, + ) + + return out, final_state + + +@input_guard +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + transpose_state_layout: bool = False, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV, K]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + transpose_state_layout (bool): + Whether to use transposed state layout `[V, K]` instead of `[K, V]`. Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.kda import fused_recurrent_kda + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_recurrent_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o, final_state = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + inplace_final_state=False, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + transpose_state_layout=transpose_state_layout, + ) + return o, final_state diff --git a/fla/ops/kda/gate.py b/fla/ops/kda/gate.py new file mode 100644 index 0000000000000000000000000000000000000000..112e287bc0e2a0168f260295313e59acf3755334 --- /dev/null +++ b/fla/ops/kda/gate.py @@ -0,0 +1,454 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# This file is modified and supported by the Moonshot AI Team + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +def naive_kda_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Torch reference implementation for KDA gate computation. + + Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])) + + Args: + g (torch.Tensor): + Input tensor of shape `[..., H, K]`. + A_log (torch.Tensor): + Parameter tensor with `H` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[H * K]`. + + Returns: + Output tensor of shape `[..., H, K]` . + """ + H, _ = g.shape[-2:] + g = g.float() + if dt_bias is not None: + g = g + dt_bias.view(H, -1) + + g = (-A_log.view(H, 1).float().exp() * F.softplus(g.float())).to(output_dtype) + return g + + +def naive_kda_lowerbound_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + lower_bound: float = -5.0, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H, _ = g.shape[-2:] + g = g.float() + if dt_bias is not None: + g = g + dt_bias.view(H, -1) + g = lower_bound * F.sigmoid(A_log.view(H, 1).exp() * g) + return g.to(output_dtype) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["dt_bias"] is not None, + "HAS_BETA": lambda args: args["beta"] is not None, + 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({"BT": BT}, num_warps=num_warps, num_stages=num_stages) + for BT in BT_LIST_AUTOTUNE + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3] + ], + key=["H", "D"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def kda_gate_fwd_kernel( + g, + A_log, + dt_bias, + beta, + yg, + yb, + lower_bound, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_BETA: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_yg = tl.make_block_ptr(yg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + # [BT, BD] + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) + b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) + if not USE_LOWER_BOUND: + b_yg = -exp(b_A) * softplus(b_g) + else: + b_yg = lower_bound * tl.sigmoid(exp(b_A) * b_g) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0, 1)) + + if HAS_BETA: + p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_yb = tl.make_block_ptr(yb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_yb = tl.sigmoid(tl.load(p_b, boundary_check=(0,)).to(tl.float32)) + tl.store(p_yb, b_yb.to(p_yb.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["dt_bias"] is not None, + "HAS_BETA": lambda args: args["beta"] is not None, + 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3] + ], + key=["H", "D"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def kda_gate_bwd_kernel( + g, + A_log, + dt_bias, + beta, + dyg, + dyb, + dg, + dA, + dbeta, + lower_bound, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_BETA: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_dg = tl.make_block_ptr(dg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_dyg = tl.make_block_ptr(dyg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + + # [BT, BD] + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_dyg = tl.load(p_dyg, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) + b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + # [BT, BD] + if not USE_LOWER_BOUND: + b_A = -exp(b_A) + b_yg = b_A * softplus(b_g) + b_dg = b_A * (b_dyg * tl.sigmoid(b_g)) + b_dA = tl.sum(tl.sum(b_dyg * b_yg, 1), 0) + else: + b_A = exp(b_A) + b_inner = b_A * b_g + b_sig = tl.sigmoid(b_inner) + b_dsig = b_sig * (1.0 - b_sig) + # Common term: dy * (LB * dsig) + b_d_inner_term = b_dyg * (lower_bound * b_dsig) + # dg = d_inner_term * A + b_dg = b_d_inner_term * b_A + b_dA = tl.sum(tl.sum(b_dg * b_g, 1), 0) + + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(dA + i_t * H + i_h, b_dA) + + if HAS_BETA: + p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(dbeta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dyb = tl.make_block_ptr(dyb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_db = tl.load(p_dyb, boundary_check=(0,)).to(tl.float32) * b_b * (1.0 - b_b) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def kda_gate_fwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H, K = g.shape[-2:] + T = g.numel() // (H * K) + + yg = torch.empty_like(g, dtype=output_dtype) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + beta=None, + yg=yg, + yb=None, + T=T, + H=H, + D=K, + BD=triton.next_power_of_2(K), + lower_bound=lower_bound, + ) + return yg + + +def kda_gate_bwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + dyg: torch.Tensor | None = None, + lower_bound: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + H, K = g.shape[-2:] + T = g.numel() // (H * K) + BT = 32 + NT = triton.cdiv(T, BT) + + dg = torch.empty_like(g, dtype=torch.float32) + dA = A_log.new_empty(NT, H, dtype=torch.float32) + + grid = (triton.cdiv(T, BT), H) + kda_gate_bwd_kernel[grid]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + beta=None, + dyg=dyg, + dyb=None, + dg=dg, + dA=dA, + dbeta=None, + T=T, + H=H, + D=K, + BT=BT, + BD=triton.next_power_of_2(K), + lower_bound=lower_bound, + ) + + dg = dg.view_as(g).type_as(g) + dA = dA.sum(0).view_as(A_log).type_as(A_log) + dbias = dg.view(-1, H * K).sum(0).to(dt_bias) if dt_bias is not None else None + + return dg, dA, dbias + + +class KDAGateFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + yg = kda_gate_fwd( + g=g, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + output_dtype=output_dtype + ) + ctx.save_for_backward(g, A_log, dt_bias) + ctx.lower_bound = lower_bound + return yg + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dyg: torch.Tensor): + g, A_log, dt_bias = ctx.saved_tensors + dg, dA, dbias = kda_gate_bwd( + g=g, + A_log=A_log, + dt_bias=dt_bias, + dyg=dyg, + lower_bound=ctx.lower_bound + ) + return dg, dA, dbias, None, None + + +@torch.compiler.disable +def fused_kda_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Fused KDA gate computation with autograd support. + + Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])) + + Args: + g (torch.Tensor): + Input tensor of shape `[..., H, K]`. + A_log (torch.Tensor): + Parameter tensor with `H` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[H * K]`. + + Returns: + Output tensor of shape `[..., H, K]`. + """ + return KDAGateFunction.apply(g, A_log, dt_bias, lower_bound, output_dtype) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["dt_bias"] is not None, + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BS': BS}, num_warps=num_warps) + for BS in BS_LIST + for num_warps in [2, 4, 8] + ], + key=['H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def kda_gate_chunk_cumsum_vector_kernel( + s, + A_log, + dt_bias, + o, + scale, + cu_seqlens, + chunk_indices, + lower_bound, + T, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + + # Apply dt_bias if exists + if HAS_BIAS: + p_b = tl.make_block_ptr(dt_bias + i_h * S, (S,), (1,), (i_s * BS,), (BS,), (0,)) + b_bias = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_s = b_s + b_bias[None, :] + + b_A = tl.load(A_log + i_h).to(tl.float32) + if not USE_LOWER_BOUND: + # Apply gate: -exp(A_log) * softplus(g + bias) + b_gate = -exp(b_A) * softplus(b_s) + else: + b_gate = lower_bound * tl.sigmoid(exp(b_A) * b_s) + + # Apply chunk local cumsum + if REVERSE: + b_o = tl.cumsum(b_gate, axis=0, reverse=True) + else: + b_o = tl.cumsum(b_gate, axis=0) + + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@input_guard +def kda_gate_chunk_cumsum( + g: torch.Tensor, + A_log: torch.Tensor, + chunk_size: int, + scale: float = None, + dt_bias: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, + lower_bound: float | None = None, + **kwargs, +) -> torch.Tensor: + if cu_seqlens is not None: + assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" + assert len(g.shape) == 4 + B, T, H, S = g.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=g_org, + A_log=A_log, + dt_bias=dt_bias, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + T=T, + H=H, + S=S, + BT=BT, + REVERSE=False, + ) + return g diff --git a/fla/ops/kda/naive.py b/fla/ops/kda/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..d7bcf45b47bffe13500c80e9a7ea7744e947135a --- /dev/null +++ b/fla/ops/kda/naive.py @@ -0,0 +1,100 @@ + + +import torch +from einops import rearrange + + +def naive_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, +): + dtype = v.dtype + B, T, H, K, V = *q.shape, v.shape[-1] + if scale is None: + scale = K ** -0.5 + + q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta]) + q = q * scale + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + for i in range(0, T): + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] + S = S * g_i[..., None].exp() + S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2)) + o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S) + if not output_final_state: + S = None + return o.to(dtype), S + + +def naive_chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, +): + dtype = v.dtype + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + NT = T // BT + if scale is None: + scale = K ** -0.5 + assert T % BT == 0 + + q, k, v, g, beta = map(lambda x: rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float), [q, k, v, g, beta]) + q = q * scale + g = g.cumsum(-2) + + # note that diagonal is masked. + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + + A = torch.zeros(*q.shape[:-1], BT, dtype=torch.float, device=q.device) + for i in range(BT): + k_i = k[..., i, :] + g_i = g[..., i:i+1, :] + A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i) + A = A * beta[..., None] + + A = -A.masked_fill(mask, 0) + for i in range(1, BT): + A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2) + A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta[..., None, :] + + w = A @ (g.exp() * k) + u = A @ v + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, NT): + # [B, H, BT, ...] + q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i] + A = torch.zeros(B, H, BT, BT, dtype=torch.float, device=q.device) + for j in range(BT): + k_j = k[:, :, i, j] + g_j = g[:, :, i, j:j+1, :] + A[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j) + A = A.masked_fill(mask, 0) + v_i = u_i - w_i @ S + o[:, :, i] = (q_i * g_i.exp()) @ S + A @ v_i + S = S * rearrange(g_i[:, :, -1].exp(), 'b h k -> b h k 1') + S += rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, 'b h c k -> b h k c') @ v_i + if not output_final_state: + S = None + return rearrange(o, 'b h n c d -> b (n c) h d').to(dtype), S diff --git a/fla/ops/kda/wy_fast.py b/fla/ops/kda/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..27c6c2e4c3eed87c7fe525bd966a1aeec118da5f --- /dev/null +++ b/fla/ops/kda/wy_fast.py @@ -0,0 +1,311 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2 +from fla.utils import autotune_cache_kwargs, check_shared_mem + + +@triton.heuristics({ + 'STORE_QG': lambda args: args['qg'] is not None, + 'STORE_KG': lambda args: args['kg'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kda_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_qg = tl.make_block_ptr(qg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.).to(tl.float32) + b_kg = b_k * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) + p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kda_kernel( + k, + v, + beta, + gk, + A, + dA, + dw, + du, + dk, + dk2, + dv, + db, + dg, + dg2, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_b = tl.load(p_b, boundary_check=(0,)) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2 + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg2 = tl.make_block_ptr(dg2 + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk_exp = exp2(tl.load(p_gk, boundary_check=(0, 1))) + b_kbg = b_k * b_b[:, None] * b_gk_exp + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + b_dk = b_dkbg * b_gk_exp * b_b[:, None] + tl.load(p_dk, boundary_check=(0, 1)) + b_db += tl.sum(b_dkbg * b_k * b_gk_exp, 1) + b_dg = b_kbg * b_dkbg + tl.load(p_dg, boundary_check=(0, 1)) + + tl.store(p_dk2, b_dk.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg2, b_dg.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + b_dA = tl.where(m_A, -b_dA, 0) + + # if using gk, save dA first and handle dk in another kernel + p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + qg = torch.empty_like(q) if q is not None else None + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kda_kernel[(NT, B*H)]( + q=q, + k=k, + qg=qg, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u, qg, kg + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + gk: torch.Tensor, + A: torch.Tensor, + dk: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk2 = torch.empty_like(dk, dtype=torch.float) + dv = torch.empty_like(v) + dg2 = torch.empty_like(gk, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + prepare_wy_repr_bwd_kda_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + gk=gk, + A=A, + dA=dA, + dw=dw, + du=du, + dk=dk, + dk2=dk2, + dv=dv, + db=db, + dg=dg, + dg2=dg2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dk = dk2 + dg = dg2 + return dk, dv, db, dg, dA diff --git a/fla/ops/lightning_attn/__init__.py b/fla/ops/lightning_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f2684e74721530e6cd4a72714e657d2b3e5188a2 --- /dev/null +++ b/fla/ops/lightning_attn/__init__.py @@ -0,0 +1,8 @@ + +from .chunk import chunk_lightning_attn +from .fused_recurrent import fused_recurrent_lightning_attn + +__all__ = [ + 'chunk_lightning_attn', + 'fused_recurrent_lightning_attn', +] diff --git a/fla/ops/lightning_attn/chunk.py b/fla/ops/lightning_attn/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..674d53b6dab600e0f8e5e634dcea904179e5ccf5 --- /dev/null +++ b/fla/ops/lightning_attn/chunk.py @@ -0,0 +1,82 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.simple_gla.chunk import chunk_simple_gla + + +@torch.compiler.disable +def chunk_lightning_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_idx: int, + num_layers: int, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + layer_idx (int): + The index of the current layer. + num_layers (int): + The total number of layers. Both `layer_idx` and `num_layers` are used to compute the decay factor. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + + H = q.shape[2] + g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) + return chunk_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + g_gamma=g_gamma, + initial_state=initial_state, + output_final_state=output_final_state, + head_first=head_first, + cu_seqlens=cu_seqlens, + ) diff --git a/fla/ops/lightning_attn/fused_recurrent.py b/fla/ops/lightning_attn/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..05182e9073de3a1f5a63d2bb15e39872fe9b2cf1 --- /dev/null +++ b/fla/ops/lightning_attn/fused_recurrent.py @@ -0,0 +1,82 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla + + +def fused_recurrent_lightning_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_idx: int, + num_layers: int, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + layer_idx (int): + The index of the current layer. + num_layers (int): + The total number of layers. Both `layer_idx` and `num_layers` are used to compute the decay factor. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + H = q.shape[2] + g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) + return fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + head_first=head_first, + ) diff --git a/fla/ops/linear_attn/__init__.py b/fla/ops/linear_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..56d329a7c987a17913cf83bd16560265d4f34dc2 --- /dev/null +++ b/fla/ops/linear_attn/__init__.py @@ -0,0 +1,10 @@ + +from .chunk import chunk_linear_attn +from .fused_chunk import fused_chunk_linear_attn +from .fused_recurrent import fused_recurrent_linear_attn + +__all__ = [ + 'chunk_linear_attn', + 'fused_chunk_linear_attn', + 'fused_recurrent_linear_attn', +] diff --git a/fla/ops/linear_attn/chunk.py b/fla/ops/linear_attn/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..735e2241d3496feb4f0839064838c82b4d8890de --- /dev/null +++ b/fla/ops/linear_attn/chunk.py @@ -0,0 +1,74 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.linear_attn.utils import normalize_output +from fla.ops.simple_gla import chunk_simple_gla + + +@torch.compiler.disable +def chunk_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + normalize: bool = True, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + scale (Optional[float]): + Scale factor for the linear attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[B, H, K, V]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. + normalize (bool): + Whether to normalize the output. Default: `True`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None`. + """ + + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first: + if q.shape[1] < q.shape[2]: + raise DeprecationWarning( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = chunk_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + ) + if normalize: + o = normalize_output(q * scale, k, o) + return o, final_state diff --git a/fla/ops/linear_attn/fused_chunk.py b/fla/ops/linear_attn/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..6806f2003b5080c377194a9c258e1d54e1326b72 --- /dev/null +++ b/fla/ops/linear_attn/fused_chunk.py @@ -0,0 +1,59 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.linear_attn.utils import normalize_output +from fla.ops.simple_gla import fused_chunk_simple_gla + + +@torch.compiler.disable +def fused_chunk_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + normalize: bool = True, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + scale (Optional[float]): + Scale factor for linear attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[B, H, K, V]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. + normalize (bool): + Whether to normalize the output. Default: `True`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None` + """ + o, final_state = fused_chunk_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + if normalize: + o = normalize_output(q * scale, k, o) + return o, final_state diff --git a/fla/ops/linear_attn/fused_recurrent.py b/fla/ops/linear_attn/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..f745affa594fbe65fb4419012ab71df3450fab33 --- /dev/null +++ b/fla/ops/linear_attn/fused_recurrent.py @@ -0,0 +1,33 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.linear_attn.utils import normalize_output +from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla + + +def fused_recurrent_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + normalize: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + o, final_state = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + if normalize: + o = normalize_output(q * scale, k, o) + return o, final_state diff --git a/fla/ops/linear_attn/naive.py b/fla/ops/linear_attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf6ed46d6f20059ce9f8bf9e76804065b3059ab --- /dev/null +++ b/fla/ops/linear_attn/naive.py @@ -0,0 +1,62 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +from einops import rearrange + +from fla.ops.linear_attn.utils import normalize_output + + +def naive_recurrent_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float | None = None, + normalize: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + dtype = q.dtype + if scale is None: + scale = q.shape[-1] ** -0.5 + B, T, H, K, V = *q.shape, v.shape[-1] + q, k, v = map(lambda x: x.to(torch.float32), (q, k, v)) + o = torch.empty_like(v) + + S = torch.zeros((B, H, K, V), device=q.device, dtype=torch.float32) + if initial_state is not None: + S = S + initial_state + for t in range(T): + S = S + torch.einsum('b h k, b h v -> b h k v', k[:, t], v[:, t]) + o[:, t] = torch.einsum('b h k v, b h k -> b h v', S, q[:, t] * scale) + if normalize: + o = normalize_output(q * scale, k, o) + return o.to(dtype), S if output_final_state else None + + +def naive_chunk_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + normalize: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + if scale is None: + scale = q.shape[-1] ** -0.5 + chunk_size = 64 + q = rearrange(q, 'b (n c) h d -> b h n c d', c=chunk_size) * scale + k = rearrange(k, 'b (n c) h d -> b h n c d', c=chunk_size) + v = rearrange(v, 'b (n c) h d -> b h n c d', c=chunk_size) + kv = k.transpose(-1, -2) @ v + kv = kv.cumsum(2) + kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) + inter = q @ kv + intra = (( + q @ k.transpose(-1, -2)).masked_fill_( + torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), + 0, + )) @ v + o = inter + intra + if normalize: + o = normalize_output(q * scale, k, o) + return rearrange(o, 'b h n c d -> b (n c) h d') diff --git a/fla/ops/linear_attn/utils.py b/fla/ops/linear_attn/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..66ca8bfed8cc715987808b185217074cddae96b2 --- /dev/null +++ b/fla/ops/linear_attn/utils.py @@ -0,0 +1,8 @@ + +import torch + + +def normalize_output(q: torch.Tensor, k: torch.Tensor, o: torch.Tensor) -> torch.Tensor: + k = k.cumsum(1) + z = (q * k).sum(-1, keepdim=True) + return o / (z + 1e-10) diff --git a/fla/ops/log_linear_attn/__init__.py b/fla/ops/log_linear_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a4fb56280686ea42fee81767cd450fd21947a25b --- /dev/null +++ b/fla/ops/log_linear_attn/__init__.py @@ -0,0 +1,6 @@ + +from .chunk import chunk_log_linear_attn + +__all__ = [ + 'chunk_log_linear_attn', +] diff --git a/fla/ops/log_linear_attn/chunk.py b/fla/ops/log_linear_attn/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..83dc3d414ec779615c0783e43f4674ecc13633d1 --- /dev/null +++ b/fla/ops/log_linear_attn/chunk.py @@ -0,0 +1,1917 @@ +import math +import warnings +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import reduce + +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.op import exp +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + +BLOCK_K = 64 + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + }, +) +@triton.autotune( + configs=[ + triton.Config({"BK": BLOCK_K}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_fwd_kernel( + q, + k, + v, + g, + level_scales, + llut, + o, + h0, + ht, + offsets, + new_offsets, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + L_IN: tl.constexpr, + L_OUT: tl.constexpr, + MIN_LEVEL: tl.constexpr, + MAX_LEVEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, +): + p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) + b_llut = tl.load(p_llut, boundary_check=(0, 1)) + # parallel over sequences and heads + i_k = tl.program_id(0) + i_nh = tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_i = tl.arange(0, BT) + + # For hierarchical masking + num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + i_idx = o_i[:, None] # BT x 1 + j_idx = o_i[None, :] # 1 x BT + + # This is not great. + # See issue: https://github.com/triton-lang/triton/discussions/1313 + KV_0_CREATED = MIN_LEVEL <= 1 and MAX_LEVEL >= 0 + KV_1_CREATED = MIN_LEVEL <= 2 and MAX_LEVEL >= 0 + KV_2_CREATED = MIN_LEVEL <= 3 and MAX_LEVEL >= 1 + KV_3_CREATED = MIN_LEVEL <= 4 and MAX_LEVEL >= 2 + KV_4_CREATED = MIN_LEVEL <= 5 and MAX_LEVEL >= 3 + KV_5_CREATED = MIN_LEVEL <= 6 and MAX_LEVEL >= 4 + KV_6_CREATED = MIN_LEVEL <= 7 and MAX_LEVEL >= 5 + KV_7_CREATED = MIN_LEVEL <= 8 and MAX_LEVEL >= 6 + KV_8_CREATED = MIN_LEVEL <= 9 and MAX_LEVEL >= 7 + KV_9_CREATED = MIN_LEVEL <= 10 and MAX_LEVEL >= 8 + KV_10_CREATED = MIN_LEVEL <= 11 and MAX_LEVEL >= 9 + KV_11_CREATED = MIN_LEVEL <= 12 and MAX_LEVEL >= 10 + + kv_0 = tl.zeros([BK, V], dtype=tl.float32) + kv_1 = tl.zeros([BK, V], dtype=tl.float32) + kv_2 = tl.zeros([BK, V], dtype=tl.float32) + kv_3 = tl.zeros([BK, V], dtype=tl.float32) + kv_4 = tl.zeros([BK, V], dtype=tl.float32) + kv_5 = tl.zeros([BK, V], dtype=tl.float32) + kv_6 = tl.zeros([BK, V], dtype=tl.float32) + kv_7 = tl.zeros([BK, V], dtype=tl.float32) + kv_8 = tl.zeros([BK, V], dtype=tl.float32) + kv_9 = tl.zeros([BK, V], dtype=tl.float32) + kv_10 = tl.zeros([BK, V], dtype=tl.float32) + kv_11 = tl.zeros([BK, V], dtype=tl.float32) + + offset = 0 # total number to cached tokens + first_chunk_index = 0 # next chunk index to compute + if USE_INITIAL_STATE: + offset = tl.load(offsets + i_n) + + first_chunk_index = offset // BT + + if KV_0_CREATED and (first_chunk_index & 1 > 0): + p_kv_0 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 0) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_0 = tl.load(p_kv_0, boundary_check=(0, 1)) + if KV_1_CREATED and (first_chunk_index & 2 > 0): + p_kv_1 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 1) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_1 = tl.load(p_kv_1, boundary_check=(0, 1)) + if KV_2_CREATED and (first_chunk_index & 4 > 0): + p_kv_2 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 2) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_2 = tl.load(p_kv_2, boundary_check=(0, 1)) + if KV_3_CREATED and (first_chunk_index & 8 > 0): + p_kv_3 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 3) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_3 = tl.load(p_kv_3, boundary_check=(0, 1)) + if KV_4_CREATED and (first_chunk_index & 16 > 0): + p_kv_4 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 4) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_4 = tl.load(p_kv_4, boundary_check=(0, 1)) + if KV_5_CREATED and (first_chunk_index & 32 > 0): + p_kv_5 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 5) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_5 = tl.load(p_kv_5, boundary_check=(0, 1)) + if KV_6_CREATED and (first_chunk_index & 64 > 0): + p_kv_6 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 6) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_6 = tl.load(p_kv_6, boundary_check=(0, 1)) + if KV_7_CREATED and (first_chunk_index & 128 > 0): + p_kv_7 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 7) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_7 = tl.load(p_kv_7, boundary_check=(0, 1)) + if KV_8_CREATED and (first_chunk_index & 256 > 0): + p_kv_8 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 8) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_8 = tl.load(p_kv_8, boundary_check=(0, 1)) + if KV_9_CREATED and (first_chunk_index & 512 > 0): + p_kv_9 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 9) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_9 = tl.load(p_kv_9, boundary_check=(0, 1)) + if KV_10_CREATED and (first_chunk_index & 1024 > 0): + p_kv_10 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 10) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_10 = tl.load(p_kv_10, boundary_check=(0, 1)) + if KV_11_CREATED and (first_chunk_index & 2048 > 0): + p_kv_11 = tl.make_block_ptr( + h0 + ((i_n * L_IN + 11) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + kv_11 = tl.load(p_kv_11, boundary_check=(0, 1)) + + NT = tl.cdiv(T, BT) + output_offset = -1 * (offset % BT) + for i_t in range(NT): + b_h_ptrs = level_scales + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut + b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) + + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_q = tl.make_block_ptr( + q + bos * K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + p_k = tl.make_block_ptr( + k + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), + ) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, 0), + (BT, V), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + ((bos * H + i_h) * (K // BK) + i_k) * V, + (T, V), + (H * (K // BK) * V, 1), + (i_t * BT + output_offset, 0), + (BT, V), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0,)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + m_t = i_t * BT + o_i < T + b_s = (tl.dot(b_q, b_k) * tl.where((i_idx >= j_idx) & m_t[:, None] & m_t[None, :], tl.exp(b_g[:, None] - b_g[None, :]), 0)).to( + b_q.dtype, + ) * b_h + + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_o = tl.zeros((BT, V), dtype=tl.float32) + if MIN_LEVEL == 0: + b_o += tl.dot(b_s, b_v) + + chunk_index = ( + first_chunk_index + i_t + ) # index of the chunk over the entire sequence, including the offset + + if MIN_LEVEL <= 0 and MAX_LEVEL >= 0: + if chunk_index & 1: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_0.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 1 and MAX_LEVEL >= 1: + if chunk_index & 2: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 1), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_1.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 2 and MAX_LEVEL >= 2: + if chunk_index & 4: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 2), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_2.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 3 and MAX_LEVEL >= 3: + if chunk_index & 8: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 3), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_3.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 4 and MAX_LEVEL >= 4: + if chunk_index & 16: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 4), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_4.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 5 and MAX_LEVEL >= 5: + if chunk_index & 32: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 5), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_5.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 6 and MAX_LEVEL >= 6: + if chunk_index & 64: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 6), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_6.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 7 and MAX_LEVEL >= 7: + if chunk_index & 128: # 8192 - 16384 + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 7), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_7.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 8 and MAX_LEVEL >= 8: + if chunk_index & 256: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 8), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_8.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 9 and MAX_LEVEL >= 9: + if chunk_index & 512: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 9), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_9.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 10 and MAX_LEVEL >= 10: + if chunk_index & 1024: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 10), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_10.to(b_q.dtype)) * tl.exp(b_g)[:, None] + if MIN_LEVEL <= 11 and MAX_LEVEL >= 11: + if chunk_index & 2048: + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, num_intra_levels + 11), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0, 1)) + b_o += tl.dot((b_l * b_q), kv_11.to(b_q.dtype)) * tl.exp(b_g)[:, None] + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + if i_t < NT - 1 or T % BT == 0: + # Only apply the state update if the last chunk is a full chunk. + # Otherwise, it needs to be included in the next kernel call. + + # update the recurrent states + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + if KV_0_CREATED: + kv_0 *= tl.exp(b_g_last) + if KV_1_CREATED: + kv_1 *= tl.exp(b_g_last) + if KV_2_CREATED: + kv_2 *= tl.exp(b_g_last) + if KV_3_CREATED: + kv_3 *= tl.exp(b_g_last) + if KV_4_CREATED: + kv_4 *= tl.exp(b_g_last) + if KV_5_CREATED: + kv_5 *= tl.exp(b_g_last) + if KV_6_CREATED: + kv_6 *= tl.exp(b_g_last) + if KV_7_CREATED: + kv_7 *= tl.exp(b_g_last) + if KV_8_CREATED: + kv_8 *= tl.exp(b_g_last) + if KV_9_CREATED: + kv_9 *= tl.exp(b_g_last) + if KV_10_CREATED: + kv_10 *= tl.exp(b_g_last) + if KV_11_CREATED: + kv_11 *= tl.exp(b_g_last) + + b_v = (b_v * tl.exp(b_g_last - b_g)[:, None]).to(b_v.dtype) + if MIN_LEVEL <= 1: + kv_0 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 2: + kv_1 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 3: + kv_2 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 4: + kv_3 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 5: + kv_4 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 6: + kv_5 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 7: + kv_6 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 8: + kv_7 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 9: + kv_8 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 10: + kv_9 += tl.dot(b_k, b_v) + elif MIN_LEVEL == 11: + kv_10 += tl.dot(b_k, b_v) + + check_value = (~chunk_index & (chunk_index + 1)) - 1 + + if MIN_LEVEL <= 1 and MAX_LEVEL >= 0: + if check_value & 1: + kv_1 += kv_0 + kv_0 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 2 and MAX_LEVEL >= 1: + if check_value & 2: + kv_2 += kv_1 + kv_1 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 3 and MAX_LEVEL >= 2: + if check_value & 4: + kv_3 += kv_2 + kv_2 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 4 and MAX_LEVEL >= 3: + if check_value & 8: + kv_4 += kv_3 + kv_3 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 5 and MAX_LEVEL >= 4: + if check_value & 16: + kv_5 += kv_4 + kv_4 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 6 and MAX_LEVEL >= 5: + if check_value & 32: + kv_6 += kv_5 + kv_5 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 7 and MAX_LEVEL >= 6: + if check_value & 64: + kv_7 += kv_6 + kv_6 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 8 and MAX_LEVEL >= 7: + if check_value & 128: + kv_8 += kv_7 + kv_7 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 9 and MAX_LEVEL >= 8: + if check_value & 256: + kv_9 += kv_8 + kv_8 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 10 and MAX_LEVEL >= 9: + if check_value & 512: + kv_10 += kv_9 + kv_9 = tl.zeros([BK, V], dtype=tl.float32) + if MIN_LEVEL <= 11 and MAX_LEVEL >= 10: + if check_value & 1024: + kv_11 += kv_10 + kv_10 = tl.zeros([BK, V], dtype=tl.float32) + + chunk_index = offset // BT + T // BT + + if STORE_FINAL_STATE: + if (MIN_LEVEL <= 0 and MAX_LEVEL >= 0) and (chunk_index & 1 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 0) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_0, boundary_check=(0, 1)) + if (MIN_LEVEL <= 1 and MAX_LEVEL >= 1) and (chunk_index & 2 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 1) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_1, boundary_check=(0, 1)) + if (MIN_LEVEL <= 2 and MAX_LEVEL >= 2) and (chunk_index & 4 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 2) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_2, boundary_check=(0, 1)) + if (MIN_LEVEL <= 3 and MAX_LEVEL >= 3) and (chunk_index & 8 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 3) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_3, boundary_check=(0, 1)) + if (MIN_LEVEL <= 4 and MAX_LEVEL >= 4) and (chunk_index & 16 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 4) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_4, boundary_check=(0, 1)) + if (MIN_LEVEL <= 5 and MAX_LEVEL >= 5) and (chunk_index & 32 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 5) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_5, boundary_check=(0, 1)) + if (MIN_LEVEL <= 6 and MAX_LEVEL >= 6) and (chunk_index & 64 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 6) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_6, boundary_check=(0, 1)) + if (MIN_LEVEL <= 7 and MAX_LEVEL >= 7) and (chunk_index & 128 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 7) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_7, boundary_check=(0, 1)) + if (MIN_LEVEL <= 8 and MAX_LEVEL >= 8) and (chunk_index & 256 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 8) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_8, boundary_check=(0, 1)) + if (MIN_LEVEL <= 9 and MAX_LEVEL >= 9) and (chunk_index & 512 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 9) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_9, boundary_check=(0, 1)) + if (MIN_LEVEL <= 10 and MAX_LEVEL >= 10) and (chunk_index & 1024 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 10) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_10, boundary_check=(0, 1)) + if (MIN_LEVEL <= 11 and MAX_LEVEL >= 11) and (chunk_index & 2048 > 0): + p_kv = tl.make_block_ptr( + ht + ((i_n * L_OUT + 11) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + tl.store(p_kv, kv_11, boundary_check=(0, 1)) + + tl.store(new_offsets + i_n, (offset // BT) * BT + T) + + +@triton.heuristics({ + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +}) +@triton.jit(do_not_specialize=["T"]) +def copy_input_kernel( + q, + k, + v, + g, + level_scales, + cu_seqlens, + q_prev, + k_prev, + v_prev, + g_prev, + level_scales_prev, + offsets, + q_new, + k_new, + v_new, + g_new, + level_scales_new, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # parallel over sequences and heads + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + offset = tl.load(offsets + i_n) + input_offset = -1 * (offset % BT) + + NT = tl.cdiv(T, BT) + + for i_t in range(NT): + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT + input_offset,), (BT,), (0,), + ) + p_q = tl.make_block_ptr( + q + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), + ) + p_k = tl.make_block_ptr( + k + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), + ) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT + input_offset, 0), + (BT, V), + (1, 0), + ) + p_g_new = tl.make_block_ptr( + g_new + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), + ) + p_q_new = tl.make_block_ptr( + q_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_k_new = tl.make_block_ptr( + k_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_v_new = tl.make_block_ptr( + v_new + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, 0), + (BT, V), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0,)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + if i_t == 0: + p_g_prev = tl.make_block_ptr( + g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), + ) + p_q_prev = tl.make_block_ptr( + q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), + ) + p_k_prev = tl.make_block_ptr( + k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), + ) + p_v_prev = tl.make_block_ptr( + v_prev + (i_n * BT * H + i_h) * V, + (BT, V), + (H * V, 1), + (0, 0), + (BT, V), + (1, 0), + ) + + b_g += tl.load(p_g_prev, boundary_check=(0,)) + b_q += tl.load(p_q_prev, boundary_check=(0, 1)) + b_k += tl.load(p_k_prev, boundary_check=(0, 1)) + b_v += tl.load(p_v_prev, boundary_check=(0, 1)) + + tl.store(p_g_new, b_g, boundary_check=(0,)) + tl.store(p_q_new, b_q, boundary_check=(0, 1)) + tl.store(p_k_new, b_k, boundary_check=(0, 1)) + tl.store(p_v_new, b_v, boundary_check=(0, 1)) + + for i in range(L): + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT + input_offset, i), + (BT, 1), + (1, 0), + ) + p_l_new = tl.make_block_ptr( + level_scales_new + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (i_t * BT, i), + (BT, 1), + (1, 0), + ) + b_l = tl.load(p_l, boundary_check=(0,)) + if i_t == 0: + p_l_prev = tl.make_block_ptr( + level_scales_prev + (i_n * BT * H + i_h) * L, + (BT, L), + (H * L, 1), + (0, i), + (BT, 1), + (1, 0), + ) + b_l += tl.load(p_l_prev, boundary_check=(0,)) + tl.store(p_l_new, b_l, boundary_check=(0,)) + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + }, +) +@triton.jit(do_not_specialize=["T"]) +def copy_last_chunk_kernel( + q, + k, + v, + g, + level_scales, + cu_seqlens, + q_prev, + k_prev, + v_prev, + g_prev, + level_scales_prev, + offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # parallel over sequences and heads + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + seq_offset = (T // BT) * BT + + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (seq_offset,), (BT,), (0,)) + p_q = tl.make_block_ptr( + q + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), + ) + p_k = tl.make_block_ptr( + k + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), + ) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (seq_offset, 0), + (BT, V), + (1, 0), + ) + p_g_prev = tl.make_block_ptr( + g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), + ) + p_q_prev = tl.make_block_ptr( + q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), + ) + p_k_prev = tl.make_block_ptr( + k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), + ) + p_v_prev = tl.make_block_ptr( + v_prev + (i_n * BT * H + i_h) * V, (BT, V), (H * V, 1), (0, 0), (BT, V), (1, 0), + ) + + tl.store(p_g_prev, tl.load(p_g, boundary_check=(0,)), boundary_check=(0,)) + tl.store(p_q_prev, tl.load(p_q, boundary_check=(0, 1)), boundary_check=(0, 1)) + tl.store(p_k_prev, tl.load(p_k, boundary_check=(0, 1)), boundary_check=(0, 1)) + tl.store(p_v_prev, tl.load(p_v, boundary_check=(0, 1)), boundary_check=(0, 1)) + + for i in range(L): + p_l = tl.make_block_ptr( + level_scales + (bos * H + i_h) * L, + (T, L), + (H * L, 1), + (seq_offset, i), + (BT, 1), + (1, 0), + ) + p_l_prev = tl.make_block_ptr( + level_scales_prev + (i_n * BT * H + i_h) * L, + (BT, L), + (H * L, 1), + (0, i), + (BT, 1), + (1, 0), + ) + tl.store(p_l_prev, tl.load(p_l, boundary_check=(0,)), boundary_check=(0,)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + restore_value=["dh", "dg_last"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_bwd_kernel_dhg( + do, + q, + g, + l, + h_l, + dh, + dg_last, + ell, + T, + cu_seqlens, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + NT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # parallel over batches and heads + i_k = tl.program_id(0) + i_nh = tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + b_dh = tl.zeros([BK, V], dtype=tl.float32) + + num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + + for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): + p_dh = tl.make_block_ptr( + dh + ((i_n * NT + i_t) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + b_dh_old = tl.load(p_dh, boundary_check=(0, 1)) + + if (i_t & (1 << ell)) == 0: # store the chunk + tl.store( + p_dh, b_dh.to(p_dh.dtype.element_ty) + b_dh_old, boundary_check=(0, 1), + ) + # if you are about the transition to compute, reset to zeros + if i_t > 0 and ((i_t - 1) & (1 << ell)) > 0: + b_dh = tl.zeros([BK, V], dtype=tl.float32) + if i_t & (1 << ell): + p_h = tl.make_block_ptr( + h_l + ((i_n * NT + i_t) * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, 0), + (BK, V), + (1, 0), + ) + + b_h = tl.load(p_h, boundary_check=(0, 1)) + p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h + tl.atomic_add(p_dg_last, tl.sum(b_h * (b_dh + b_dh_old))) + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.exp(tl.load(g + bos * H + last_idx * H + i_h)) + b_dh *= b_g_last + if i_t & (1 << ell): # compute this chunk + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), + ) + p_q = tl.make_block_ptr( + q + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), + ) + p_do = tl.make_block_ptr( + do + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, 0), + (BT, V), + (1, 0), + ) + p_l = tl.make_block_ptr( + l + (bos * H + i_h) * L + num_intra_levels + ell, + (T,), + (H * L,), + (i_t * BT,), + (BT,), + (0,), + ) + b_l = tl.load(p_l, boundary_check=(0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * (tl.exp(b_g) * b_l)[None, :]).to(b_q.dtype) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + b_s = tl.dot(b_q, b_do).to(b_q.dtype) + b_dh += b_s + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + restore_value=["dq", "dg"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_bwd_kernel_hdqgl( + do, + q, + k, + v, + g, + l, + h_l, + dq, + dg, + dl, + ell, + T, + cu_seqlens, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + NT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # parallel over batches and heads + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + b_h = tl.zeros([V, K], dtype=tl.float32) + + num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + + for i_t in range(tl.cdiv(T, BT)): + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + if i_t & (1 << ell): # compute and store derivatives + p_do = tl.make_block_ptr( + do + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, 0), + (BT, V), + (1, 0), + ) + p_q = tl.make_block_ptr( + q + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_l = tl.make_block_ptr( + l + (bos * H + i_h) * L + num_intra_levels + ell, + (T,), + (H * L,), + (i_t * BT,), + (BT,), + (0,), + ) + p_dq = tl.make_block_ptr( + dq + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, 0), + (BT, K), + (1, 0), + ) + p_dg = tl.make_block_ptr( + dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), + ) + p_dl = tl.make_block_ptr( + dl + (bos * H + i_h) * L + num_intra_levels + ell, + (T,), + (H * L,), + (i_t * BT,), + (BT,), + (0,), + ) + p_h = tl.make_block_ptr( + h_l + ((i_n * NT + i_t) * H + i_h) * K * V, + (V, K), + (1, V), + (0, 0), + (V, K), + (0, 1), + ) + + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_l = tl.load(p_l, boundary_check=(0,)) + + b_dlq = tl.exp(b_g)[:, None] * tl.dot(b_do, b_h.to(b_do.dtype)) + + b_dl = tl.sum(b_dlq * b_q, axis=1) + b_dg = b_l * b_dl + + tl.store(p_h, b_h, boundary_check=(0, 1)) + b_dq_old = tl.load(p_dq, boundary_check=(0, 1)) + tl.store( + p_dq, + (b_l[:, None] * b_dlq).to(p_dq.dtype.element_ty) + b_dq_old, + boundary_check=(0, 1), + ) + tl.store(p_dl, b_dl.to(p_dl.dtype.element_ty), boundary_check=(0,)) + b_dg_old = tl.load(p_dg, boundary_check=(0,)) + tl.store( + p_dg, b_dg.to(p_dg.dtype.element_ty) + b_dg_old, boundary_check=(0,), + ) + if ((i_t + 1) & (1 << ell)) == 0: + b_h = tl.zeros([V, K], dtype=tl.float32) + + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + b_h *= tl.exp(b_g_last) + if (i_t & (1 << ell)) == 0: # update the state + p_k = tl.make_block_ptr( + k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (V, T), + (1, H * V), + (0, i_t * BT), + (V, BT), + (0, 1), + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_k = (b_k * tl.exp(b_g_last - b_g)[:, None]).to(b_k.dtype) + b_h += tl.dot(b_v, b_k) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + restore_value=["dk", "dg", "dg_last"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_bwd_kernel_dkg( + dh, + k, + v, + g, + dg_last, + dk, + dg, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + NT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_i = tl.arange(0, BT) + o_t = i_t * BT + o_i + m_t = o_t < T + + p_dh = tl.make_block_ptr( + dh + ((i_n * NT + i_t) * H + i_h) * K * V, + (V, K), + (1, V), + (0, 0), + (V, K), + (0, 1), + ) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, 0), + (BT, V), + (1, 0), + ) + p_dk = tl.make_block_ptr( + dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h + b_dg_last = tl.load(p_dg_last) + + b_dg_last *= tl.exp(b_g_last) + b_dk = tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] * tl.dot(b_v, b_dh).to(b_v.dtype) + b_dg = tl.load(p_dg, boundary_check=(0,)) + b_dg -= tl.sum(b_k * b_dk, axis=1) + b_dg_last += tl.sum(b_dk * b_k) + + b_dg = tl.where(o_i < BT - 1, b_dg, b_dg + b_dg_last) + + tl.store(p_dg, b_dg, boundary_check=(0,)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + restore_value=["dv"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_bwd_kernel_dv( + dh, + k, + g, + dv, + T, + cu_seqlens, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + NT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_dh = tl.make_block_ptr( + dh + ((i_n * NT + i_t) * H + i_h) * K * V, + (K, V), + (V, 1), + (0, 0), + (K, V), + (1, 0), + ) + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + p_dv = tl.make_block_ptr( + dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), + ) + + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dv = tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] * tl.dot(b_k, b_dh).to(b_k.dtype) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [4] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V"], + restore_value=["dl", "dq", "dk", "dv", "dg"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunkwise_bwd_kernel_diag( + do, + q, + k, + v, + g, + l, + llut, + mask, + dq, + dk, + dv, + dg, + dl, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) + b_llut = tl.load(p_llut, boundary_check=(0, 1)) + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_i = tl.arange(0, BT) + i_idx = o_i[:, None] # BT x 1 + j_idx = o_i[None, :] # 1 x BT + + b_h_ptrs = l + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut + b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) + + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_q = tl.make_block_ptr(q + bos * K, (K, T), (1, K), (0, i_t * BT), (K, BT), (0, 1)) + p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, (V, T), (1, H * V), (0, i_t * BT), (V, BT), (0, 1), + ) + p_do = tl.make_block_ptr( + do + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), + ) + p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dq = tl.make_block_ptr( + dq + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_dk = tl.make_block_ptr( + dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), + ) + p_dv = tl.make_block_ptr( + dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0,)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dq = tl.load(p_dq, boundary_check=(0, 1)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dg = tl.load(p_dg, boundary_check=(0,)) + + b_s = (tl.dot(b_k, b_q)).to(b_q.dtype) + # Apply causal and padding masks + m_t = i_t * BT + o_i < T + b_a = tl.where((i_idx >= j_idx) & m_t[:, None] & m_t[None, :], tl.exp(b_g[:, None] - b_g[None, :]), 0) + b_dv += tl.dot((b_s * tl.trans(b_a * b_h)).to(b_do.dtype), b_do) + b_ds = tl.dot(b_do, b_v) * b_a + b_dl = b_ds * tl.trans(b_s) + b_dg += tl.sum(b_dl * b_h, axis=1) + b_dg -= tl.sum(b_dl * b_h, axis=0) + b_ds = (b_ds * b_h).to(b_k.dtype) + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.trans(tl.dot(b_q, b_ds)) + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + + for i in range(num_intra_levels): + p_mask = tl.make_block_ptr(mask + i * (BT * BT), (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) + b_mask = tl.load(p_mask, boundary_check=(0, 1)) + dl_i = tl.sum(tl.where(b_mask == 1, b_dl, 0), axis=1) + p_dl_i = tl.make_block_ptr(dl + (bos * H + i_h) * L + i, (T,), (H * L,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dl_i, dl_i, boundary_check=(0,)) + + +def construct_binary_level_mask(level, T): + if level == 0: + return torch.diag(torch.ones(T, dtype=torch.bool)) + + indices = torch.cartesian_prod(torch.arange(T), torch.arange(T)) + + mask = torch.where( + torch.logical_and( + torch.logical_and( + indices[:, 0] % (1 << level) >= (1 << (level - 1)), + indices[:, 1] + (1 << (level - 1)) + >= indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), + ), + indices[:, 1] < indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), + ).view(T, T), + 1, + 0, + ) + + return mask + + +def level_lut(BT, device): + lut = torch.zeros((BT, BT), dtype=torch.int32, device=device) + for level in range(1, ceil_log(BT, 2) + 1): + mask = construct_binary_level_mask(level, BT).to(device) + lut = torch.where(mask.to(torch.bool), level, lut) + return lut + + +def masks(BT, device): + masks = [] + for level in range(0, ceil_log(BT, 2) + 1): + mask = construct_binary_level_mask(level, BT).to(device).to(torch.int32) + masks.append(mask) + return torch.stack(masks) + + +def ceil_div(x: int, y: int) -> int: + return math.ceil(x / y) + + +def ceil_log(x: int, b: int) -> int: + return math.ceil(math.log(x, b)) + + +@dataclass +class LogLinearAttentionState: + ht: torch.Tensor + offsets: torch.Tensor + q_prev: torch.Tensor + k_prev: torch.Tensor + v_prev: torch.Tensor + g_prev: torch.Tensor + level_scales_prev: torch.Tensor + + +class ChunkLogLinearAttentionFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + g, + level_scales, + initial_state, + output_final_state, + cu_seqlens, + ): + B, T, G, K = k.shape + _, _, H, V = v.shape + _, _, _, L = level_scales.shape + + if G != 1: + raise ValueError("Group dimension must be 1.") + + if not math.log2(V).is_integer(): + raise ValueError( + "Head dimension must be a power of two. Please pad the head dimension to the next power of two.", + ) + + if K % BLOCK_K != 0: + raise ValueError(f"State dimension must be divisible by {BLOCK_K}.") + + if triton.__version__ > "3.2.0": + warnings.warn("Triton>3.2.0 detected, which is known to have worse performance. " + "For optimal performance, it is recommended to install Triton==3.2.0 (if possible).") + + BT = 64 # chunk size + + h0 = initial_state.ht if initial_state is not None else None + offsets = initial_state.offsets if initial_state is not None else None + + if cu_seqlens is None: + NT = ceil_div(T + (torch.max(offsets) if offsets is not None else 0), BT) + MAX_LEVEL = ceil_log(NT, 2) - 1 + else: + NT = max( + [ + ceil_div( + cu_seqlens[i + 1] + - cu_seqlens[i] + + (offsets[i] if offsets is not None else 0), + BT, + ) + for i in range(len(cu_seqlens) - 1) + ], + ) + MAX_LEVEL = ceil_log(NT, 2) - 1 + B = len(cu_seqlens) - 1 + + if MAX_LEVEL > 10: + raise ValueError("Sequence length must be less than 2**17") + + S0 = B if cu_seqlens is None else 1 + o = torch.zeros( + (S0, T, H, (K // BLOCK_K), V), + dtype=v.dtype, + device=v.device, + ) + + if initial_state is not None: + if cu_seqlens is not None: + cu_seqlens = cu_seqlens + F.pad(torch.cumsum(offsets % BT), (1, 0)) + else: + assert (offsets == offsets[0]).all() + T += offsets[0].item() % BT + S1 = cu_seqlens[-1] if cu_seqlens is not None else T + q_new = torch.zeros((S0, S1, G, K), dtype=q.dtype, device=q.device) + k_new = torch.zeros((S0, S1, G, K), dtype=k.dtype, device=k.device) + v_new = torch.zeros((S0, S1, H, V), dtype=v.dtype, device=v.device) + g_new = torch.zeros((S0, S1, H), dtype=g.dtype, device=g.device) + level_scales_new = torch.zeros((S0, S1, H, L), dtype=level_scales.dtype, device=level_scales.device) + + copy_input_kernel[(B * H,)]( + q=q, + k=k, + v=v, + g=g, + level_scales=level_scales, + cu_seqlens=cu_seqlens, + q_prev=initial_state.q_prev, + k_prev=initial_state.k_prev, + v_prev=initial_state.v_prev, + g_prev=initial_state.g_prev, + level_scales_prev=initial_state.level_scales_prev, + q_new=q_new, + k_new=k_new, + v_new=v_new, + g_new=g_new, + level_scales_new=level_scales_new, + offsets=offsets, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + ) + q = q_new + k = k_new + v = v_new + g = g_new + level_scales = level_scales_new + + # Store one extra level (MAX_LEVEL + 2) in case the length is multiple of 2 + ht = ( + torch.zeros((B, MAX_LEVEL + 2, H, K, V), dtype=torch.float, device=v.device) + if output_final_state + else None + ) + + new_offsets = torch.zeros((B,), dtype=torch.int32, device=v.device) + g = chunk_local_cumsum(g, chunk_size=BT, cu_seqlens=cu_seqlens) + + def grid(meta): + return (triton.cdiv(K, meta["BK"]), B * H) + + l_in = h0.shape[1] if initial_state is not None else None + l_out = ht.shape[1] if output_final_state else None + + ctx.llut = level_lut(BT, v.device) + + chunkwise_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + level_scales=level_scales, + llut=ctx.llut, + o=o, + h0=h0, + ht=ht, + offsets=offsets, + new_offsets=new_offsets, + cu_seqlens=cu_seqlens, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + L_IN=l_in, + L_OUT=l_out, + MIN_LEVEL=0, + MAX_LEVEL=MAX_LEVEL, + ) + + ctx.save_for_backward(q, k, v, g, level_scales, initial_state, cu_seqlens) + ctx.chunk_size = BT + + if output_final_state: + q_prev = torch.zeros((B, BT, G, K), dtype=q.dtype, device=q.device) + k_prev = torch.zeros((B, BT, G, K), dtype=k.dtype, device=k.device) + v_prev = torch.zeros((B, BT, H, V), dtype=v.dtype, device=v.device) + g_prev = torch.zeros((B, BT, H), dtype=g.dtype, device=g.device) + level_scales_prev = torch.zeros((B, BT, H, L), dtype=level_scales.dtype, device=level_scales.device) + + copy_last_chunk_kernel[(B * H,)]( + q=q, + k=k, + v=v, + g=g, + level_scales=level_scales, + cu_seqlens=cu_seqlens, + q_prev=q_prev, + k_prev=k_prev, + v_prev=v_prev, + g_prev=g_prev, + level_scales_prev=level_scales_prev, + offsets=new_offsets, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + ) + + final_state = LogLinearAttentionState( + ht=ht, + offsets=new_offsets, + q_prev=q_prev, + k_prev=k_prev, + v_prev=v_prev, + g_prev=g_prev, + level_scales_prev=level_scales_prev, + ) + return o.sum(dim=-2), final_state + + return o.sum(dim=-2), None + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + if triton.__version__ < "3.1.0": + raise ValueError("Triton>=3.1.0 is required") + + q, k, v, g, level_scales, initial_state, cu_seqlens = ctx.saved_tensors + chunk_size = ctx.chunk_size + llut = ctx.llut + mask = masks(chunk_size, v.device) + + if initial_state is not None: + raise NotImplementedError( + "Backward pass is not implemented for log-linear attention with a prefilled kernel.", + ) + + B, T, G, K = k.shape + assert G == 1, "Multi-head attention is not supported" + _, _, H, V = v.shape + _, _, _, L = level_scales.shape + BT = chunk_size + if cu_seqlens is not None: + NT = max( + [ + ceil_div(cu_seqlens[i + 1] - cu_seqlens[i], BT) + for i in range(len(cu_seqlens) - 1) + ], + ) + else: + NT = ceil_div(T, BT) + + if cu_seqlens is not None: + B = len(cu_seqlens) - 1 + + dh = torch.zeros((B, NT, H, K, V), dtype=v.dtype, device=v.device) + dq = torch.zeros((B if cu_seqlens is None else 1, T, H, K), dtype=v.dtype, device=v.device) + dk = torch.zeros((B if cu_seqlens is None else 1, T, H, K), dtype=v.dtype, device=v.device) + dv = torch.zeros_like(v) + dg = torch.zeros(g.shape, dtype=torch.float, device=v.device) + dl = torch.zeros(level_scales.shape, dtype=torch.float, device=v.device) + h_l = torch.zeros((B, NT, H, K, V), dtype=torch.float, device=v.device) + dg_last = torch.zeros((B, NT, H), dtype=torch.float, device=v.device) + do = do.to(v.dtype) + + grid = (B * H,) + + def grid_f(meta): + return (triton.cdiv(K, meta["BK"]), B * H) + + grid_t = (NT, B * H) + + num_inter_chunk_levels = ceil_log(NT, 2) + for ell in range(num_inter_chunk_levels - 1, -1, -1): + chunkwise_bwd_kernel_hdqgl[grid]( + do=do, + q=q, + k=k, + v=v, + g=g, + l=level_scales, + h_l=h_l, + dq=dq, + dg=dg, + dl=dl, + cu_seqlens=cu_seqlens, + ell=ell, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + NT=NT, + ) + chunkwise_bwd_kernel_dhg[grid_f]( + do=do, + q=q, + g=g, + l=level_scales, + h_l=h_l, + dh=dh, + dg_last=dg_last, + cu_seqlens=cu_seqlens, + ell=ell, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + NT=NT, + ) + + chunkwise_bwd_kernel_dkg[grid_t]( + dh=dh, + k=k, + v=v, + g=g, + dg_last=dg_last, + dk=dk, + dg=dg, + cu_seqlens=cu_seqlens, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + NT=NT, + ) + + chunkwise_bwd_kernel_dv[grid_t]( + dh=dh, + k=k, + g=g, + dv=dv, + cu_seqlens=cu_seqlens, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + NT=NT, + ) + + chunkwise_bwd_kernel_diag[grid_t]( + do=do, + q=q, + k=k, + v=v, + g=g, + l=level_scales, + llut=llut, + mask=mask, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + dl=dl, + cu_seqlens=cu_seqlens, + T=T, + H=H, + K=K, + V=V, + L=L, + BT=BT, + ) + + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens).to(g.dtype) + + dq = reduce(dq, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) + dk = reduce(dk, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) + return dq, dk, dv, dg, dl, None, None, None + + +@torch.compiler.disable +def chunk_log_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + level_scales: torch.Tensor, + initial_state: LogLinearAttentionState | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + level_scales (torch.Tensor): + Scales for each level of shape `[B, T, H, L]`. + initial_state (Optional[LogLinearAttentionState]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of type `LogLinearAttentionState` if `output_final_state=True` else `None`. + + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + + o, final_state = ChunkLogLinearAttentionFunction.apply( + q, + k, + v, + g, + level_scales, + initial_state, + output_final_state, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/log_linear_attn/naive.py b/fla/ops/log_linear_attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..aa16b814e0fb5f096583121481a81be6df5eba4c --- /dev/null +++ b/fla/ops/log_linear_attn/naive.py @@ -0,0 +1,50 @@ +import numpy as np +import torch + + +def segsum(x): + T = x.size(-1) + x_cumsum = torch.cumsum(x, dim=-1) + x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :] + mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool)) + x_segsum = x_segsum.masked_fill(~mask, -torch.inf) + return x_segsum + + +def construct_level_mask(level, L): + T = L.size(-1) + if level == 0: + return torch.diag_embed(L[..., level, :]) + + indices = torch.cartesian_prod(torch.arange(T), torch.arange(T)).to(L.device) + + mask = torch.where( + torch.logical_and( + torch.logical_and( + indices[:, 0] % (1 << level) >= (1 << (level - 1)), + indices[:, 1] + (1 << (level - 1)) + >= indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), + ), + indices[:, 1] < indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), + ).view(T, T), + L[..., level, :].unsqueeze(-1).expand(*([-1] * (len(L.shape) - 2)), T, T), + 0, + ) + + return mask + + +def construct_H_matrix(a, L): + T = a.size(-1) + A = torch.exp(segsum(a)) + H = torch.zeros_like(A) + for level in range(int(np.ceil(np.log2(T))) + 1): + mask = construct_level_mask(level, L) + H += A * mask + return H + + +def naive_log_linear_attn(q, k, v, g, level_scales): + H = construct_H_matrix(g.permute(0, 2, 1), level_scales.permute(0, 2, 3, 1)) + M = torch.einsum("bhlc,blhn,bchn->bhlc", H, q, k) + return torch.einsum("bhlc,bchp->blhp", M, v) diff --git a/fla/ops/mesa_net/__init__.py b/fla/ops/mesa_net/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..24036866a902a98c9f5db16b6758694ea133be1f --- /dev/null +++ b/fla/ops/mesa_net/__init__.py @@ -0,0 +1,5 @@ +from .chunk import chunk_mesa_net +from .decoding_one_step import mesa_net_decoding_one_step +from .naive import naive_mesa_net_decoding_one_step, naive_mesa_net_exact + +__all__ = ['chunk_mesa_net', 'naive_mesa_net_exact', 'mesa_net_decoding_one_step', 'naive_mesa_net_decoding_one_step'] diff --git a/fla/ops/mesa_net/chunk.py b/fla/ops/mesa_net/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..f7cce869457bb028232d19ae8f12539c65688c56 --- /dev/null +++ b/fla/ops/mesa_net/chunk.py @@ -0,0 +1,382 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_h import chunk_bwd_dh +from fla.ops.mesa_net.chunk_cg_solver_bwd import chunk_mesa_cg_bwd +from fla.ops.mesa_net.chunk_cg_solver_fwd import chunk_mesa_cg_fwd +from fla.ops.mesa_net.chunk_h_fwd import chunk_mesa_fwd_h +from fla.ops.mesa_net.chunk_h_kk_intra_bwd import chunk_mesa_net_h_kk_bwd_intra_fn +from fla.ops.mesa_net.chunk_h_kv_intra_bwd import chunk_mesa_net_h_kv_bwd_intra_fn +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_fwd_mesa_net_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + lamb: torch.Tensor, + cu_seqlens: torch.Tensor, + max_CG_iteration: int = 30, + chunk_size: int = 64, + h_kk_init: torch.Tensor | None = None, + h_kv_init: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) if g is not None else None + h_kk, h_kv, h_kk_final, h_kv_final = chunk_mesa_fwd_h( + k=k, + v=v, + g=g, + beta=beta, + h_init=h_kk_init, + h_kv_init=h_kv_init, + output_final_state=output_final_state, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + q_star, o = chunk_mesa_cg_fwd( + q=q, + k=k, + h=h_kk, + h_kv=h_kv, + v=v, + g_local_cumsum=g, + beta=beta, + lamb=lamb, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + max_CG_iteration=max_CG_iteration, + chunk_indices=chunk_indices, + ) + return g, q_star, o, (h_kk_final, h_kv_final) + + +def chunk_fwd_mesa_net_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + lamb: torch.Tensor, + q_star: torch.Tensor, # should be cached in the forward pass + do: torch.Tensor, + cu_seqlens: torch.Tensor, + max_CG_iteration: int = 30, + chunk_size: int = 64, + h_kk_init: torch.Tensor | None = None, + h_kv_init: torch.Tensor | None = None, + dh_kv_final: torch.Tensor | None = None, + dh_kk_final: torch.Tensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + # recompute the hidden states, which is quite cheap + h_kk, h_kv, _, _ = chunk_mesa_fwd_h( + k=k, + v=v, + g=g, + beta=beta, + h_init=h_kk_init, + h_kv_init=h_kv_init, + output_final_state=False, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + dh_kv, dh0_kv = chunk_bwd_dh( + q=q_star, + k=k, + v=v, + g=g, + gk=None, + gv=None, + do=do, + h0=h_kv_init, + dht=dh_kv_final, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + scale=1, + ) + dq, dk_beta, dv, dg = chunk_mesa_net_h_kv_bwd_intra_fn( + q_star=q_star, + k=k, + v=v, + beta=beta, + h_kv=h_kv, + dh_kv=dh_kv, + g=g, + do=do, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dq = chunk_mesa_cg_bwd( + dq=dq, + k=k, + h=h_kk, + g_local_cumsum=g, + beta=beta, + lamb=lamb, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + max_CG_iteration=max_CG_iteration, + output_dtype=torch.float16, + chunk_indices=chunk_indices, + ) + dh_kk, dh0_kk = chunk_bwd_dh( + q=dq, + k=k, + v=k, + g=g, + gk=None, + gv=None, + do=q_star, + h0=h_kk_init, + dht=-dh_kk_final if dh_kk_final is not None else None, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + scale=1, + ) + dk, dg2, dlamb, dbeta = chunk_mesa_net_h_kk_bwd_intra_fn( + k=k, + g=g, + beta=beta, + h=h_kk, + dh=dh_kk, + dk_beta=dk_beta, + q_star=q_star, + dq=dq, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dg.add_(dg2) + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens).to(g) + return dq, dk, dv, dg, dbeta, dlamb, -dh0_kk if dh0_kk is not None else None, dh0_kv if dh0_kv is not None else None + + +class ChunkMesaNetFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + g, + beta, + lamb, + cu_seqlens, + cu_seqlens_cpu, + max_CG_iteration, + h_kk_init, + h_kv_init, + output_final_state, + use_qk_l2norm_in_kernel, + ): + chunk_size = 64 + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q, output_dtype=torch.float16) + k, k_rstd = l2norm_fwd(k, output_dtype=torch.float16) + else: + q_rstd, k_rstd = None, None + q = q.to(torch.float16) + k = k.to(torch.float16) + + g_cumsum, q_star, o, (h_kk_final, h_kv_final) = chunk_fwd_mesa_net_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + lamb=lamb, + cu_seqlens=cu_seqlens, + max_CG_iteration=max_CG_iteration, + chunk_size=chunk_size, + h_kk_init=h_kk_init, + h_kv_init=h_kv_init, + output_final_state=output_final_state, + chunk_indices=chunk_indices, + ) + ctx.max_CG_iteration = max_CG_iteration + ctx.chunk_size = chunk_size + ctx.cu_seqlens = cu_seqlens + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g_cumsum, beta, lamb, h_kk_init, h_kv_init, q_star, o, chunk_indices) + return o, h_kk_final, h_kv_final + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dh_kk_final=None, dh_kv_final=None): + q, q_rstd, k, k_rstd, v, g, beta, lamb, h_kk_init, h_kv_init, q_star, o, chunk_indices = ctx.saved_tensors + + max_CG_iteration = ctx.max_CG_iteration + chunk_size = ctx.chunk_size + cu_seqlens = ctx.cu_seqlens + dq, dk, dv, dg, dbeta, dlamb, dh0_kk, dh0_kv = chunk_fwd_mesa_net_bwd( + q=q, k=k, v=v, g=g, beta=beta, lamb=lamb, q_star=q_star, do=do, + cu_seqlens=cu_seqlens, max_CG_iteration=max_CG_iteration, chunk_size=chunk_size, + h_kk_init=h_kk_init, h_kv_init=h_kv_init, dh_kv_final=dh_kv_final, dh_kk_final=dh_kk_final, + chunk_indices=chunk_indices, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq, dk, dv.to(v), dg.to(g), dbeta.to(beta), dlamb.to(lamb), None, None, None, dh0_kk, dh0_kv, None, None + + +@torch.compiler.disable +def chunk_mesa_net( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + lamb: torch.Tensor, + h_kk_init: torch.Tensor | None = None, + h_kv_init: torch.Tensor | None = None, + output_final_state: bool = False, + max_CG_iteration: int = 30, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]` + k (torch.Tensor): + keys of shape `[B, T, H, K]`. Should be l2-normalized before passing in. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + decay factors of shape `[B, T, H]`. Note that `g` should be in log space, that is, `g = log(decay_factor) < 0`. + Recommended input dtype: `torch.float32`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. Recommended input dtype: `torch.float32`. + lamb (torch.Tensor): + lambdas of shape `[B, T, H]`. Recommended input dtype: `torch.float32`. + h_kk_init (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + h_kv_init (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + max_CG_iteration (int): + Maximum number of conjugate gradient iterations for solving the linear system. Default: `30`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Do l2 normalization on Q and K in the kernel for saving GPU memory. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + (final_states_kk, final_states_kv) (Tuple[torch.Tensor, torch.Tensor]): + Final states of shape `[N, H, K, K]` and `[N, H, K, V]` if `output_final_state=True` else `(None, None)`. + Recall that MesaNet has two states, `h_kk` and `h_kv`! + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.mesa_net import chunk_mesa_net + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 16, 128, 128 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, dtype=torch.float32, device='cuda')) + >>> beta = torch.rand(B, T, H, dtype=torch.float32, device='cuda').sigmoid() + # lower bound is 0.25 for numerical stability + >>> lamb = F.softplus(torch.rand(H, K, dtype=torch.float32, device='cuda')) + 0.25 + >>> init_state_kk = torch.randn(B, H, K, V, dtype=torch.float32, device='cuda') + >>> init_state_kv = torch.randn(B, H, K, V, dtype=torch.float32, device='cuda') + >>> o, (final_state_kk, final_state_kv) = chunk_mesa_net( + q, k, v, beta, lamb, + h_kk_init=init_state_kk, + h_kv_init=init_state_kv, + max_CG_iteration=30, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, (final_state_kk_var, final_state_kv_var) = chunk_mesa_net( + q, k, v, beta, lamb, + h_kk_init=init_state_kk, + h_kv_init=init_state_kv, + max_CG_iteration=30, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + B, T, H, K = q.shape + assert k.shape == (B, T, H, K), "k must be of shape (batch size, seq len, num head, head dim)." + assert v.shape == (B, T, H, K), "v must be of shape (batch size, seq len, num head, head dim)." + assert g.shape == (B, T, H), "g must be of shape (batch size, seq len, num head)." + assert beta.shape == (B, T, H), "beta must be of shape (batch size, seq len, num head)." + assert lamb.shape == (H, K), "lamb must be of shape (num head, key dim)." + + if h_kv_init is not None: + assert h_kv_init.dtype == torch.float32, "h_kv_init must be in float32." + if cu_seqlens is None: + assert h_kv_init.shape == (B, H, K, K), "h_kv_init must be of shape (batch size, num head, head dim, head dim)." + if h_kk_init is not None: + assert h_kk_init.dtype == torch.float32, "h_kk_init must be in float32." + if cu_seqlens is None: + assert h_kk_init.shape == (B, H, K, K), "h_kk_init must be of shape (batch size, num head, head dim, head dim)." + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if h_kk_init is not None and h_kk_init.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {h_kk_init.shape[0]}.", + ) + if h_kv_init is not None and h_kv_init.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {h_kv_init.shape[0]}.", + ) + o, final_state_kk, final_state_kv = ChunkMesaNetFunction.apply( + q, + k, + v, + g, + beta, + lamb, + cu_seqlens, + cu_seqlens_cpu, + max_CG_iteration, + h_kk_init, + h_kv_init, + output_final_state, + use_qk_l2norm_in_kernel, + ) + return o, final_state_kk, final_state_kv diff --git a/fla/ops/mesa_net/chunk_cg_solver_bwd.py b/fla/ops/mesa_net/chunk_cg_solver_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5c2a71e5242a054931b80f84b280d2869a7271 --- /dev/null +++ b/fla/ops/mesa_net/chunk_cg_solver_bwd.py @@ -0,0 +1,162 @@ + +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp + + +@triton.jit() +def chunk_update_once( + b_p, + b_k, + b_v, + b_m, + b_g_exp_q, + b_h, + b_lamb, +): + b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) + b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) + if b_lamb is not None: + b_o += b_lamb[None, :] * b_p + return b_o + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_mesa_cg_dim64_kernel( + dq, + dq_final, + k, + h, + g, + beta, + lamb, + cu_seqlens, + chunk_indices, + T, + max_CG_iteration: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + dq += (bos * H + i_h) * K + dq_final += (bos * H + i_h) * K + k += (bos * H + i_h) * K + h += (i_tg * H + i_h).to(tl.int64) * K * K + + g += bos * H + i_h + beta += bos * H + i_h + lamb += i_h * K + + p_q = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) + b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) + + b_m = exp(b_g[:, None] - b_g[None, :]) * b_beta[None, :] + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) + b_g_exp_q = tl.exp(b_g)[:, None] + + b_x = tl.zeros([BT, BK], dtype=tl.float32) + b_p = tl.zeros([BT, BK], dtype=tl.float32) + b_r = tl.zeros([BT, BK], dtype=tl.float32) + + b_x += b_q * 0. + b_r += b_q + b_p += b_q + b_delta_old = tl.sum(b_r*b_r, axis=1) + for _ in range(max_CG_iteration): + b_o = chunk_update_once(b_p, b_k, b_k, b_m, b_g_exp_q, b_h, b_lamb) + alpha = b_delta_old / (tl.sum(b_p*b_o, axis=1) + 1e-5) + b_x += alpha[:, None] * b_p + b_r = b_r - alpha[:, None] * b_o + b_delta_new = tl.sum(b_r*b_r, axis=1) + b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p + b_delta_old = b_delta_new + + p_q_final = tl.make_block_ptr(dq_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_mesa_cg_bwd( + dq: torch.Tensor, + k: torch.Tensor, + h: torch.Tensor, + g_local_cumsum: torch.Tensor, + beta: torch.Tensor, + lamb: torch.Tensor, # lambda + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + max_CG_iteration: int = 30, + output_dtype: torch.dtype | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K = dq.shape + assert K <= 128, "head dimension must be less than 128" + assert chunk_size <= 64 or K <= 64, "either chunk size or head dimension must be no greater than 64" + dq_final = torch.empty_like(dq, dtype=dq.dtype if output_dtype is None else output_dtype) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + BK = max(triton.next_power_of_2(K), 16) + grid = (NT, H*B) + + chunk_fwd_mesa_cg_dim64_kernel[grid]( + dq=dq, + dq_final=dq_final, + k=k, + h=h, + g=g_local_cumsum, + beta=beta, + lamb=lamb, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + max_CG_iteration=max_CG_iteration, + T=T, + H=H, + K=K, + BT=chunk_size, + BK=BK, + num_warps=4, + num_stages=1, + ) + return dq_final diff --git a/fla/ops/mesa_net/chunk_cg_solver_fwd.py b/fla/ops/mesa_net/chunk_cg_solver_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..ad0a748327e3b5028b1a4711890ac1622e16274c --- /dev/null +++ b/fla/ops/mesa_net/chunk_cg_solver_fwd.py @@ -0,0 +1,186 @@ + +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp + + +@triton.jit() +def chunk_update_once( + b_p, + b_k, + b_v, + b_m, + b_g_exp_q, + b_h, + b_lamb, +): + b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) + b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) + if b_lamb is not None: + b_o += b_lamb[None, :] * b_p + return b_o + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_fwd_mesa_cg_dim64_kernel( + q, + q_final, + k, + h, + o, + v, + h_kv, + g, + beta, + lamb, + cu_seqlens, + chunk_indices, + T, + max_CG_iteration: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + q += (bos * H + i_h) * K + q_final += (bos * H + i_h) * K + k += (bos * H + i_h) * K + h += (i_tg * H + i_h).to(tl.int64) * K * K + g += bos * H + i_h + beta += bos * H + i_h + lamb += i_h * K + + o += (bos * H + i_h) * K + v += (bos * H + i_h) * K + h_kv += (i_tg * H + i_h).to(tl.int64) * K * K + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) + + b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) + + b_m = exp(b_g[:, None] - b_g[None, :]) * b_beta[None, :] + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) + b_g_exp_q = tl.exp(b_g)[:, None] + + b_x = tl.zeros([BT, BK], dtype=tl.float32) + b_p = tl.zeros([BT, BK], dtype=tl.float32) + b_r = tl.zeros([BT, BK], dtype=tl.float32) + + b_x += b_q * 0. + b_r += b_q + b_p += b_r + b_delta_old = tl.sum(b_r*b_r, axis=1) + for i in range(max_CG_iteration): + b_o = chunk_update_once(b_p, b_k, b_k, b_m, b_g_exp_q, b_h, b_lamb) + alpha = b_delta_old / (tl.sum(b_p*b_o, axis=1) + 1e-5) + b_x += alpha[:, None] * b_p + b_r = b_r - alpha[:, None] * b_o + b_delta_new = tl.sum(b_r*b_r, axis=1) + b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p + b_delta_old = b_delta_new + + p_q_final = tl.make_block_ptr(q_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) + + p_h_kv = tl.make_block_ptr(h_kv, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + b_h_kv = tl.load(p_h_kv, boundary_check=(0, 1)) + p_v = tl.make_block_ptr(v, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_o = chunk_update_once(b_x, b_k, b_v, b_m, b_g_exp_q, b_h_kv, None) + p_o = tl.make_block_ptr(o, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_mesa_cg_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + h_kv: torch.Tensor, + g_local_cumsum: torch.Tensor, + beta: torch.Tensor, + lamb: torch.Tensor, # lambda + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + max_CG_iteration: int = 30, + output_dtype: torch.dtype | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K = q.shape + assert K <= 128, "head dimension must be less than 128" + assert chunk_size <= 64 or K <= 64, "either chunk size or head dimension must be no greater than 64" + q_final = torch.empty_like(q, dtype=q.dtype if output_dtype is None else output_dtype) + + assert v is not None, "v must be provided if calculate_output is True" + assert h_kv is not None, "h_kv must be provided if calculate_output is True" + o = torch.empty_like(v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + BK = max(triton.next_power_of_2(K), 16) + grid = (NT, H*B) + + chunk_fwd_mesa_cg_dim64_kernel[grid]( + q=q, + q_final=q_final, + o=o, + v=v, + h_kv=h_kv, + k=k, + h=h, + g=g_local_cumsum, + beta=beta, + lamb=lamb, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + max_CG_iteration=max_CG_iteration, + T=T, + H=H, + K=K, + BT=chunk_size, + BK=BK, + num_warps=4, + num_stages=1, + ) + return q_final, o diff --git a/fla/ops/mesa_net/chunk_h_fwd.py b/fla/ops/mesa_net/chunk_h_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..119f86c00db4e393f97a08baafca011c39234bbb --- /dev/null +++ b/fla/ops/mesa_net/chunk_h_fwd.py @@ -0,0 +1,170 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h_init'] is not None, + 'STORE_FINAL_STATE': lambda args: args['h_final'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_mesa_net_fwd_kernel_h( + k, + v, + beta, + g, + h, + h_kv, + h_init, + h_kv_init, + h_final, + h_kv_final, + cu_seqlens, + split_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + NS = tl.cdiv(T, BS) + boh = tl.load(split_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + NS = tl.cdiv(T, BS) + boh = i_n * NS + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + b_h_kv = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h_kv0 = tl.make_block_ptr(h_kv_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h_kv = tl.load(p_h_kv0, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT): + i_s = i_t // (BS // BT) + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT, ), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V + p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_h_kv = tl.make_block_ptr(h_kv + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + if i_t % (BS // BT) == 0: + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_h_kv, b_h_kv.to(p_h_kv.dtype.element_ty), boundary_check=(0, 1)) + + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + last_idx = min((i_t + 1) * BT, T) - 1 + + # scalar decay + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h + b_h *= exp(b_g_last) + b_h_kv *= exp(b_g_last) + b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_k_decay = ((b_k * exp(b_g_last - b_g)[:, None]) * b_beta[:, None]).to(b_k2.dtype) + b_h += tl.dot(tl.trans(b_k_decay), b_k2) + b_h_kv += tl.dot(tl.trans(b_k_decay), b_v.to(b_k2.dtype)) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(h_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_h_kv_final = tl.make_block_ptr(h_kv_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h_kv_final, b_h_kv.to(p_h_kv_final.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_mesa_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + h_init: torch.Tensor, + h_kv_init: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + split_size: int | None = None, + states_in_fp32: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + assert K == V, "K must be equal to V for now" + BT = chunk_size + BS = BT if split_size is None else split_size + assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NS, split_offsets = B, triton.cdiv(T, BS), None + else: + split_offsets = prepare_chunk_offsets(cu_seqlens, BS) + N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() + + h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + h_kv = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + h_final = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + h_kv_final = k.new_empty(N, H, K, V, dtype=torch.float) + + def grid(meta): return (triton.cdiv(K, 64), triton.cdiv(V, 64), N * H) + + chunk_mesa_net_fwd_kernel_h[grid]( + k=k, + v=v, + beta=beta, + g=g, + h=h, + h_kv=h_kv, + h_init=h_init, + h_kv_init=h_kv_init, + h_final=h_final, + h_kv_final=h_kv_final, + cu_seqlens=cu_seqlens, + split_offsets=split_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + BK=64, + BV=64, + ) + return h, h_kv, h_final, h_kv_final diff --git a/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py b/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..822d977a1c6f7136185a1c0f4954d322292bd9fb --- /dev/null +++ b/fla/ops/mesa_net/chunk_h_kk_intra_bwd.py @@ -0,0 +1,191 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_mesa_net_h_kk_bwd_intra_kernel( + k, + beta, + h, + dh, + g, + q_star, + dq, + dk, + dg, + dbeta, + dk_beta, + dlamb, + cu_seqlens, + chunk_indices, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + q_star += (bos * H + i_h) * V + dq += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K*V + dh += (i_tg * H + i_h).to(tl.int64) * K*V + k += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dk_beta += (bos * H + i_h) * K + dlamb += (i_tg * H + i_h).to(tl.int64) * K + beta += (bos * H + i_h) + dbeta += (bos * H + i_h) + g += bos * H + i_h + dg += bos * H + i_h + + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dv = tl.zeros([BT, BK], dtype=tl.float32) + b_dbeta = tl.zeros([BT], dtype=tl.float32) + b_dg_last = tl.zeros([1], dtype=tl.float32) + b_dg = tl.zeros([BT], dtype=tl.float32) + + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + b_gk = tl.where(m_t, exp(b_g_last - b_g), 0) + + p_q_star = tl.make_block_ptr(q_star, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_q_star = tl.load(p_q_star, boundary_check=(0, 1)) + p_dq = tl.make_block_ptr(dq, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_dq = tl.load(p_dq, boundary_check=(0, 1)) + b_dlamb = -tl.sum(b_q_star * b_dq, axis=0) + p_dlamb = tl.make_block_ptr(dlamb, (K,), (1,), (0,), (BK,), (0,)) + tl.store(p_dlamb, b_dlamb.to(p_dlamb.dtype.element_ty), boundary_check=(0,)) + + p_h = tl.make_block_ptr(h, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_v = tl.load(p_k, boundary_check=(0, 1)) + b_k = (b_v * b_beta[:, None]).to(b_v.dtype) + + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp(b_g[:, None] - b_g[None, :]), 0) + b_s = tl.dot(b_q_star, tl.trans(b_k)) * b_m + b_ds = tl.dot(b_dq, tl.trans(b_v)) + b_dv += tl.dot(tl.trans(b_s.to(b_dq.dtype)), b_dq) + b_dm = b_s * b_ds + b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) + b_dg += tl.sum(b_dm, axis=1) + b_dg -= tl.sum(b_dm, axis=0) + b_ds = b_ds * b_m + b_dk += tl.dot(tl.trans(b_ds.to(b_q_star.dtype)), b_q_star) + + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dg += tl.sum(tl.dot(b_dq, tl.trans(b_h)) * tl.exp(b_g)[:, None] * b_q_star, axis=1) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dk2 = tl.dot(b_v, b_dh.to(b_v.dtype)) * b_gk[:, None] + b_dg -= tl.sum(b_dk2 * b_k, axis=1) + b_dg_last += tl.sum(b_dk2 * b_k) + b_dk += b_dk2 + b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_gk[:, None] + b_dh = b_dh * b_h + b_dg_last += tl.sum(b_dh) * exp(b_g_last) + + p_dk_beta = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dk -= tl.load(p_dk_beta, boundary_check=(0, 1)) + b_dbeta = tl.sum(b_dk * b_v, axis=1) + b_dk = b_dk * b_beta[:, None] + b_dv + b_dk = -b_dk + + b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, -b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dbeta = tl.make_block_ptr(dbeta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dbeta, -b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + + +def chunk_mesa_net_h_kk_bwd_intra_fn( + k: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + dh: torch.Tensor, + q_star: torch.Tensor, + dq: torch.Tensor, + dk_beta: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + B, T, H, K = k.shape + V = K + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + # CONST_TILING = 64 + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + dk = torch.empty_like(k) + dg = torch.empty_like(g) + dbeta = torch.empty_like(beta) + dlamb = torch.empty(B, NT, H, K, dtype=torch.float32, device=k.device) + grid = (NT, B * H) + + chunk_mesa_net_h_kk_bwd_intra_kernel[grid]( + k=k, + h=h, + dh=dh, + g=g, + q_star=q_star, + beta=beta, + dbeta=dbeta, + dq=dq, + dk=dk, + dk_beta=dk_beta, + dg=dg, + dlamb=dlamb, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dlamb = dlamb.sum([0, 1]) + return dk, dg, dlamb, dbeta diff --git a/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py b/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..b463ba89704e742a34c682daa58832f85afe2822 --- /dev/null +++ b/fla/ops/mesa_net/chunk_h_kv_intra_bwd.py @@ -0,0 +1,212 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.mesa_net.chunk_h_kv_intra_bwd_separate import chunk_mesa_net_h_kv_bwd_intra_separate_fn +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_mesa_net_h_kv_bwd_intra_kernel( + q_star, + k, + v, + beta, + h_kv, + g, + do, + dh_kv, + dq, + dk_beta, + dg, + dv, + cu_seqlens, + chunk_indices, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h_kv += (i_tg * H + i_h).to(tl.int64) * K*V + dh_kv += (i_tg * H + i_h).to(tl.int64) * K*V + q_star += (bos * H + i_h) * K + k += (bos * H + i_h) * K + beta += (bos * H + i_h) + g += bos * H + i_h + dg += bos * H + i_h + dq += (bos * H + i_h) * K + dk_beta += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_ds = tl.zeros([BT, BT], dtype=tl.float32) + b_dv = tl.zeros([BT, BK], dtype=tl.float32) + b_dg_last = tl.zeros([1], dtype=tl.float32) + b_dg = tl.zeros([BT], dtype=tl.float32) + + p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + + # calculation + b_dg_last += tl.sum(b_h * b_dh) + b_dg_last *= exp(b_g_last) + + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp(b_g[:, None] - b_g[None, :]), 0) + b_k = (b_k * b_beta[:, None]).to(b_k.dtype) + b_s = tl.dot(b_q, tl.trans(b_k)) * b_m + + b_ds = tl.dot(b_do, tl.trans(b_v)) + b_dm = b_s * b_ds + b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) + + b_dg += tl.sum(b_dm, axis=1) + b_dg -= tl.sum(b_dm, axis=0) + + b_g_exp_q = exp(b_g) + b_g_exp_k = tl.where(m_t, exp(-b_g + b_g_last), 0) + b_ds = b_ds * b_m + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) * b_g_exp_q[:, None] + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) * b_g_exp_k[:, None] + b_dg_last += tl.sum(b_dk * b_k) + b_dg -= tl.sum(b_dk * b_k, axis=1) + b_dg += tl.sum(b_dq * b_q, axis=1) + b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) + b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) + + b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +def chunk_mesa_net_h_kv_bwd_intra_fn( + q_star, + k, + v, + beta, + h_kv, + dh_kv, + g, + do, + cu_seqlens, + chunk_size=64, + chunk_indices: torch.LongTensor | None = None, +): + # share memory is not large enough for a single fused kernel + if not check_shared_mem('ampere'): + return chunk_mesa_net_h_kv_bwd_intra_separate_fn( + q_star=q_star, + k=k, + v=v, + beta=beta, + h_kv=h_kv, + dh_kv=dh_kv, + g=g, + do=do, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + dq = torch.empty_like(q_star, dtype=torch.float32) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dg = torch.empty_like(g) + grid = (NT, B * H) + chunk_mesa_net_h_kv_bwd_intra_kernel[grid]( + q_star=q_star, + k=k, + v=v, + beta=beta, + h_kv=h_kv, + g=g, + do=do, + dh_kv=dh_kv, + dq=dq, + dk_beta=dk, + dg=dg, + dv=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dq, dk, dv, dg diff --git a/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py b/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py new file mode 100644 index 0000000000000000000000000000000000000000..6088694c2b6b80f5b8ad046a1bf5f8ba8e088244 --- /dev/null +++ b/fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py @@ -0,0 +1,303 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( + q_star, + k, + v, + beta, + h_kv, + g, + do, + dh_kv, + dk_beta, + dg, + dv, + cu_seqlens, + chunk_indices, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h_kv += (i_tg * H + i_h).to(tl.int64) * K*V + dh_kv += (i_tg * H + i_h).to(tl.int64) * K*V + q_star += (bos * H + i_h) * K + k += (bos * H + i_h) * K + beta += (bos * H + i_h) + g += bos * H + i_h + dg += bos * H + i_h + dk_beta += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_ds = tl.zeros([BT, BT], dtype=tl.float32) + b_dv = tl.zeros([BT, BK], dtype=tl.float32) + b_dg_last = tl.zeros([1], dtype=tl.float32) + b_dg = tl.zeros([BT], dtype=tl.float32) + + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_g = tl.load(p_g, boundary_check=(0,)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + + # calculation + b_dg_last += tl.sum(b_h * b_dh) + b_dg_last *= exp(b_g_last) + + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp(b_g[:, None] - b_g[None, :]), 0) + b_k = (b_k * b_beta[:, None]).to(b_k.dtype) + b_s = tl.dot(b_q, tl.trans(b_k)) * b_m + b_ds = tl.dot(b_do, tl.trans(b_v)) + b_dm = b_s * b_ds + b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) + b_dg += tl.sum(b_dm, axis=1) + b_dg -= tl.sum(b_dm, axis=0) + b_g_exp_k = tl.where(m_t, exp(-b_g + b_g_last), 0) + b_ds = b_ds * b_m + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) * b_g_exp_k[:, None] + b_dg_last += tl.sum(b_dk * b_k) + b_dg -= tl.sum(b_dk * b_k, axis=1) + b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) + b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) + b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) + p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( + q_star, + k, + v, + beta, + h_kv, + g, + do, + dq, + dg_prev, + dg, + cu_seqlens, + chunk_indices, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + # offset calculation + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h_kv += (i_tg * H + i_h).to(tl.int64) * K*V + q_star += (bos * H + i_h) * K + k += (bos * H + i_h) * K + beta += (bos * H + i_h) + g += bos * H + i_h + dg_prev += bos * H + i_h + dg += bos * H + i_h + dq += (bos * H + i_h) * K + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dg_prev = tl.make_block_ptr(dg_prev, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_g = tl.load(p_g, boundary_check=(0,)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + + b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp(b_g[:, None] - b_g[None, :]), 0) + b_k = (b_k * b_beta[:, None]).to(b_k.dtype) + + b_ds = tl.dot(b_do, tl.trans(b_v)) * b_m + b_g_exp_q = exp(b_g) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype)) * b_g_exp_q[:, None] + b_dg = tl.sum(b_dq * b_q, axis=1) + tl.load(p_dg_prev, boundary_check=(0,)) + b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +def chunk_mesa_net_h_kv_bwd_intra_separate_fn( + q_star, + k, + v, + beta, + h_kv, + dh_kv, + g, + do, + cu_seqlens, + chunk_size=64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + dq = torch.empty_like(q_star, dtype=torch.float32) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dg = torch.empty_like(g) + grid = (NT, B * H) + chunk_mesa_net_h_kv_bwd_intra_kernel_dkv[grid]( + q_star=q_star, + k=k, + v=v, + beta=beta, + h_kv=h_kv, + g=g, + do=do, + dh_kv=dh_kv, + dk_beta=dk, + dg=dg, + dv=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dg_final = torch.empty_like(dg) + chunk_mesa_net_h_kv_bwd_intra_kernel_dq[grid]( + q_star=q_star, + k=k, + v=v, + beta=beta, + h_kv=h_kv, + g=g, + do=do, + dg=dg_final, + dg_prev=dg, + dq=dq, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dq, dk, dv, dg_final diff --git a/fla/ops/mesa_net/decoding_one_step.py b/fla/ops/mesa_net/decoding_one_step.py new file mode 100644 index 0000000000000000000000000000000000000000..b05700df47ccb16ce1559684d7c74d69552d8236 --- /dev/null +++ b/fla/ops/mesa_net/decoding_one_step.py @@ -0,0 +1,174 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import input_guard + + +@triton.jit +def mesa_net_decoding_one_step_kernel( + q, + k, + v, + g, + o, + lamb, + beta, + prev_h_kk, + prev_h_kv, + curr_h_kk, + curr_h_kv, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + MAX_CG_STEP: tl.constexpr, +): + i_nh = tl.program_id(0) + i_h = i_nh % H + + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + + p_q = q + i_nh * K + o_k + p_k = k + i_nh * K + o_k + p_v = v + i_nh * V + o_v + p_beta = beta + i_nh + p_g = g + i_nh + p_lamb = lamb + i_h * K + o_k + + b_g = exp(tl.load(p_g).to(tl.float32)) + b_beta = tl.load(p_beta).to(tl.float32) + + mask_k = o_k < K + mask_v = o_v < V + mask_kk = mask_k[:, None] & mask_k[None, :] + mask_kv = mask_k[:, None] & mask_v[None, :] + + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_lamb = tl.load(p_lamb, mask=mask_k, other=0).to(tl.float32) + + p_hkk_prev = prev_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :] + b_h_kk = tl.load(p_hkk_prev, mask=mask_kk, other=0).to(tl.float32) + + b_h_kk = b_h_kk * b_g + (b_k * b_beta)[:, None] * b_k[None, :] + + p_hkk_curr = curr_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :] + tl.store(p_hkk_curr, b_h_kk.to(p_hkk_curr.dtype.element_ty), mask=mask_kk) + + p_hkv_prev = prev_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h_kv = tl.load(p_hkv_prev, mask=mask_kv, other=0).to(tl.float32) + b_h_kv = b_h_kv * b_g + (b_k * b_beta)[:, None] * b_v[None, :] + p_hkv_curr = curr_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_hkv_curr, b_h_kv.to(p_hkv_curr.dtype.element_ty), mask=mask_kv) + + diag_mask = tl.arange(0, BK)[:, None] == tl.arange(0, BK)[None, :] + diag_mask = diag_mask & mask_kk + b_h_kk_diag = tl.sum(tl.where(diag_mask, b_h_kk, 0.0), axis=1) + + b_x = b_q / (b_h_kk_diag + b_lamb + 1e-5) + b_Hx = tl.sum(b_h_kk * b_x[:, None], axis=0) + b_r = b_q - b_Hx - b_lamb * b_x + b_p = tl.zeros([BK], dtype=tl.float32) + b_p += b_r + delta_old = tl.sum(b_r * b_r) + + for i_iter in range(MAX_CG_STEP): + b_Ap = tl.sum(b_h_kk * b_p[:, None], axis=0) + b_lamb * b_p + pAp = tl.sum(b_p * b_Ap) + alpha = delta_old / (pAp + 1e-5) + b_x = b_x + alpha * b_p + b_r = b_r - alpha * b_Ap + delta_new = tl.sum(b_r * b_r) + beta_cg = delta_new / (delta_old + 1e-5) + b_p = b_r + beta_cg * b_p + delta_old = delta_new + b_o = tl.sum(b_h_kv * b_x[:, None], axis=0) + p_o = o + i_nh * V + o_v + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + +@input_guard +def mesa_net_decoding_one_step( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + lamb: torch.Tensor, + beta: torch.Tensor, + prev_h_kk: torch.Tensor, + prev_h_kv: torch.Tensor, + max_CG_iteration: int = 30, +): + """ + Triton implementation of Mesa Net CG one step + + Args: + q (torch.Tensor): + query tensor [B, H, K] + k (torch.Tensor): + key tensor [B, H, K] + v (torch.Tensor): + value tensor [B, H, V] + g (torch.Tensor): + gate tensor [B, H] + lamb (torch.Tensor): + lambda tensor [H, K] + beta (torch.Tensor): + beta tensor [B, H] + prev_h_kk (torch.Tensor): + previous hidden state KK [B, H, K, K] + prev_h_kv (torch.Tensor): + previous hidden state KV [B, H, K, V] + max_CG_iteration (int): + maximum CG iterations + + Returns: + o (torch.Tensor): + output tensor [B, H, V] + h_kk_new (torch.Tensor): + updated hidden state KK [B, H, K, K] + h_kv_new (torch.Tensor): + updated hidden state KV [B, H, K, V] + """ + B, H, K, V = *q.shape, v.shape[-1] + + o = torch.empty((B, H, V), dtype=q.dtype, device=q.device) + curr_h_kk = torch.empty_like(prev_h_kk) + curr_h_kv = torch.empty_like(prev_h_kv) + + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + + assert BK <= 128 and BV <= 128, "BK and BV must be less than or equal to 128" + + grid = (B * H,) + mesa_net_decoding_one_step_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + o=o, + lamb=lamb, + beta=beta, + prev_h_kk=prev_h_kk, + prev_h_kv=prev_h_kv, + curr_h_kk=curr_h_kk, + curr_h_kv=curr_h_kv, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + MAX_CG_STEP=max_CG_iteration, + num_warps=4 if BK <= 64 else 8, + ) + return o, curr_h_kk, curr_h_kv diff --git a/fla/ops/mesa_net/naive.py b/fla/ops/mesa_net/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..9e7b0843342a682da2b9de189ab7251a4d27d7a1 --- /dev/null +++ b/fla/ops/mesa_net/naive.py @@ -0,0 +1,130 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +from einops import rearrange + + +def naive_mesa_net_decoding_one_step(q, k, v, g, lamb, beta, prev_h_kk, prev_h_kv, max_CG_iteration=30): + q = q.float().clone() + k = k.float().clone() + v = v.float().clone() + g = g.float().clone() + lamb = lamb.float().clone() + beta = beta.float().clone() + B, h, d = q.shape + k_beta = k * beta.unsqueeze(-1) + + h_kk = prev_h_kk * g.exp()[..., None, None] + k_beta.unsqueeze(-1) * k.unsqueeze(-2) + h_kv = prev_h_kv * g.exp()[..., None, None] + k_beta.unsqueeze(-1) * v.unsqueeze(-2) + diag_H = torch.diagonal(h_kk, dim1=-2, dim2=-1) + lamb = lamb.unsqueeze(0) + x = q / (diag_H + lamb) + r = q - (x.unsqueeze(-1) * h_kk).sum(-2) - (lamb * x) + p = r.clone() + delta_old = (r * r).sum(-1) + # CG iteration + for i in range(max_CG_iteration): + q = (p.unsqueeze(-1) * h_kk).sum(-2) + (lamb * p) + alpha = (delta_old / ((p * q).sum(-1) + 1e-5)) + x = x + (alpha[..., None] * p) + r = r - (alpha[..., None] * q) + delta_new = (r * r).sum(-1) + beta = delta_new / (delta_old + 1e-5) + p = r + (beta[..., None] * p) + delta_old = delta_new + o = (x.unsqueeze(-1) * h_kv).sum(-2) + return o, h_kk, h_kv + + +def naive_mesa_net_exact(q, k, v, g, lamb, beta, h_kk_init=None, h_kv_init=None): + B, L, h, d = q.shape + q = q.float() + k = k.float() + v = v.float() + g = g.float() + lamb = lamb.float() + beta = beta.float() + + h_kk = h_kk_init.clone() if h_kk_init is not None else torch.zeros(B, h, d, d, device=q.device) + h_kv = h_kv_init.clone() if h_kv_init is not None else torch.zeros(B, h, d, d, device=q.device) + + h_kk_all = torch.zeros(B, L, h, d, d, device=q.device) + h_kv_all = torch.zeros(B, L, h, d, d, device=q.device) + for i in range(L): + h_kk = h_kk * g[:, i, :, None, None].exp() + (k[:, i, :, :] * beta[:, i, :, None] + )[..., None] * k[:, i, :, None, :] + h_kv = h_kv * g[:, i, :, None, None].exp() + (k[:, i, :, :] * beta[:, i, :, None] + )[..., None] * v[:, i, :, None, :] + h_kk_all[:, i] = h_kk + h_kv_all[:, i] = h_kv + + q_star_gold = torch.linalg.solve(h_kk_all + torch.diag_embed(lamb)[None, None, ...], q) + o_gold = (q_star_gold[..., :, None] * h_kv_all).sum(-2) + return o_gold, h_kk, h_kv + + +def naive_mesa_net_CG(q, k, v, g, lamb, beta, chunk_size, max_CG_iteration=30, h_kk_init=None, h_kv_init=None): + B, L, h, d = q.shape + C = chunk_size + + def chunk_fn(x): return rearrange(x, 'b (n c) h ... -> b h n c ...', c=C).float() + + q_chunk, k_chunk, v_chunk, g_chunk, beta_chunk = map(chunk_fn, [q, k, v, g, beta]) + + g_chunk = g_chunk.cumsum(dim=-1) + + pairwise_decay = (g_chunk[..., None] - g_chunk[..., None, :]).exp().tril() * beta_chunk[..., None, :] + + num_chunks = q_chunk.shape[2] + + h_kv_all = torch.zeros(B, h, num_chunks, d, d, device=q.device) + h_kk_all = torch.zeros(B, h, num_chunks, d, d, device=q.device) + + h_kv = torch.zeros(B, h, d, d, device=q.device) + h_kk = torch.zeros(B, h, d, d, device=q.device) + + if h_kk_init is not None: + h_kk += h_kk_init + if h_kv_init is not None: + h_kv += h_kv_init + + chunk_decay_k = (g_chunk[..., -1, None] - g_chunk).exp() + chunk_decay_q = g_chunk.exp() + + k_chunk_processed = k_chunk * chunk_decay_k[..., None] * beta_chunk[..., None] + + for i in range(num_chunks): + h_kv_all[:, :, i, :, :] = h_kv + h_kk_all[:, :, i, :, :] = h_kk + + k_chunk_i = k_chunk[:, :, i, :, :] + v_chunk_i = v_chunk[:, :, i, :, :] + k_chunk_i_processed = k_chunk_processed[:, :, i, :, :] + + h_kk = h_kk * g_chunk[:, :, i, -1, None, None].exp() + (k_chunk_i_processed).transpose(-2, -1) @ k_chunk_i + h_kv = h_kv * g_chunk[:, :, i, -1, None, None].exp() + (k_chunk_i_processed).transpose(-2, -1) @ v_chunk_i + + # CG solver to approximate the matrix inverse solution. + # diag_H = torch.diagonal(h_kk_all, dim1=-2, dim2=-1) + lamb = lamb[None, :, None, None, :] + x = torch.zeros_like(q_chunk) + r = q_chunk - (x * chunk_decay_q[..., None]) @ h_kk_all - ((x @ k_chunk.transpose(-2, -1)) + * pairwise_decay) @ k_chunk - (lamb * x) + p = r.clone() + delta_old = (r * r).sum(-1) + + # CG iteration + for i in range(max_CG_iteration): + q = (p * chunk_decay_q[..., None]) @ h_kk_all + ((p @ k_chunk.transpose(-1, -2)) + * pairwise_decay) @ k_chunk + (lamb * p) + alpha = (delta_old / ((p * q).sum(-1) + 1e-5)) + x = x + (alpha[..., None] * p) + r = r - (alpha[..., None] * q) + delta_new = (r * r).sum(-1) + beta = delta_new / (delta_old + 1e-5) + p = r + (beta[..., None] * p) + delta_old = delta_new + + o = (x * chunk_decay_q[..., None]) @ h_kv_all + ((x @ k_chunk.transpose(-1, -2)) + * pairwise_decay) @ v_chunk + return rearrange(o, 'b h n c d -> b (n c) h d'), h_kk, h_kv diff --git a/fla/ops/nsa/__init__.py b/fla/ops/nsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a61b433234397578b12773138afb6d55e200e0e --- /dev/null +++ b/fla/ops/nsa/__init__.py @@ -0,0 +1,8 @@ + +from .naive import naive_nsa +from .parallel import parallel_nsa + +__all__ = [ + 'naive_nsa', + 'parallel_nsa', +] diff --git a/fla/ops/nsa/compression.py b/fla/ops/nsa/compression.py new file mode 100644 index 0000000000000000000000000000000000000000..06197f59c2a9866cf1114a95d2f652cd3204a5bf --- /dev/null +++ b/fla/ops/nsa/compression.py @@ -0,0 +1,543 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.attn.parallel import parallel_attn_bwd_preprocess +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_token_indices +from fla.ops.utils.op import exp, log +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit +def parallel_nsa_compression_fwd_kernel( + q, + k, + v, + o, + lse, + scale, + cu_seqlens, + token_indices, + chunk_offsets, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + + p_o = tl.make_block_ptr(o + (bos + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + # [G, BV] + b_o = tl.zeros([G, BV], dtype=tl.float32) + # max scores for the current block + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + # lse = log(acc) + m + b_acc = tl.zeros([G], dtype=tl.float32) + + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c, i_v * BV), (BC, BV), (1, 0)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BC] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + + # [G, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + if NC == 0: + b_lse = tl.zeros([G], dtype=tl.float32) + else: + b_o = b_o / b_acc[:, None] + b_lse = b_m + log(b_acc) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + if i_v == 0: + tl.store(lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G), b_lse.to(lse.dtype.element_ty)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_compression_bwd_kernel_dq( + q, + k, + v, + lse, + delta, + do, + dq, + scale, + cu_seqlens, + token_indices, + chunk_offsets, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + q += (bos + i_t) * HQ*K + do += (bos + i_t) * HQ*V + lse += (bos + i_t) * HQ + delta += (bos + i_t) * HQ + dq += (i_v * all + bos + i_t) * HQ*K + + p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + i_h * G + tl.arange(0, G) + p_delta = delta + i_h * G + tl.arange(0, G) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + + # [G, BK] + b_dq = tl.zeros([G, BK], dtype=tl.float32) + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (V, TC), (1, H*V), (i_v * BV, i_c), (BV, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_p = exp(b_s - b_lse[:, None]) + b_p = tl.where((o_c < NC)[None, :], b_p, 0) + + # [G, BV] @ [BV, BC] -> [G, BC] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [G, BC] @ [BC, BK] -> [G, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'TC']) +def parallel_nsa_compression_bwd_kernel_dkv( + q, + k, + v, + lse, + delta, + do, + dk, + dv, + cu_seqlens, + chunk_indices, + chunk_offsets, + scale, + T, + TC, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + all = B * TC + + if IS_VARLEN: + i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int32), tl.load(chunk_indices + i_c * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + # the number of compression representations in total + TC = tl.cdiv(T, BS) + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_v * all*H + boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) + + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + + for i in range(i_c * BC * BS, T): + o_c = i_c * BC + tl.arange(0, BC) + + p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) + p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + # [BC, G] + b_s = tl.dot(b_k, tl.trans(b_q)) + b_p = exp(b_s - b_lse[None, :]) + b_p = tl.where((i >= max(0, (o_c + 1) * BS - 1))[:, None], b_p, 0) + # [BC, G] @ [G, BV] -> [BC, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BC, BV] @ [BV, G] -> [BC, G] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BC, G] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BC, G] @ [G, BK] -> [BC, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +def parallel_nsa_compression_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + token_indices: torch.LongTensor | None = None, +): + B, T, HQ, K, V = *q.shape, v.shape[-1] + H = k.shape[2] + G = HQ // H + BC = BS = block_size + if check_shared_mem('hopper', q.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + else: + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, "The key dimension can not be larger than 256" + + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) if cu_seqlens is not None else None + + grid = (T, NV, B * H) + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + parallel_nsa_compression_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + lse=lse, + scale=scale, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV, + ) + return o, lse + + +def parallel_nsa_compression_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + block_size: int = 64, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + token_indices: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, HQ, K, V = *q.shape, v.shape[-1] + TC = k.shape[1] + H = k.shape[2] + G = HQ // H + BC = BS = block_size + BK = max(triton.next_power_of_2(K), 16) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + NV = triton.cdiv(V, BV) + if cu_seqlens is not None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(chunk_offsets, BC) + NC = len(chunk_indices) + else: + chunk_indices, chunk_offsets = None, None + NC = triton.cdiv(triton.cdiv(T, BS), BC) + + delta = parallel_attn_bwd_preprocess(o, do) + + dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) + grid = (T, NV, B * H) + parallel_nsa_compression_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dq=dq, + scale=scale, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV, + ) + dq = dq.sum(0) + + dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) + dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) + + grid = (NV, NC, B * H) + parallel_nsa_compression_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + TC=TC, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV, + ) + dk = dk.sum(0) + return dq, dk, dv + + +class ParallelNSACompressionFunction(torch.autograd.Function): + + @staticmethod + @contiguous + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + block_size, + scale, + cu_seqlens, + ): + ctx.dtype = q.dtype + + # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence + # for example, if the passed `cu_seqlens` is [0, 2, 6], + # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None + + o, lse = parallel_nsa_compression_fwd( + q=q, + k=k, + v=v, + block_size=block_size, + scale=scale, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + ) + ctx.save_for_backward(q, k, v, o, lse) + ctx.cu_seqlens = cu_seqlens + ctx.token_indices = token_indices + ctx.block_size = block_size + ctx.scale = scale + return o.to(q.dtype), lse + + @staticmethod + @contiguous + @autocast_custom_bwd + def backward(ctx, do, *args): + q, k, v, o, lse = ctx.saved_tensors + dq, dk, dv = parallel_nsa_compression_bwd( + q=q, + k=k, + v=v, + o=o, + lse=lse, + do=do, + block_size=ctx.block_size, + scale=ctx.scale, + cu_seqlens=ctx.cu_seqlens, + token_indices=ctx.token_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), None, None, None + + +def parallel_nsa_compression( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int = 64, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + return ParallelNSACompressionFunction.apply( + q, + k, + v, + block_size, + scale, + cu_seqlens, + ) diff --git a/fla/ops/nsa/naive.py b/fla/ops/nsa/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..e879ab8d03988169ec2952ca030f270e7857995f --- /dev/null +++ b/fla/ops/nsa/naive.py @@ -0,0 +1,101 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +from einops import repeat + + +def naive_nsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_indices: torch.LongTensor, + block_size: int = 64, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]`.. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + block_indices (torch.LongTensor): + Block indices of shape `[B, T, H, S]` if `head_first=False` else `[B, H, T, S]`. + `S` is the number of selected blocks for each query token, which is set to 16 in the paper. + block_size (int): + Selected block size. Default: 64. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + """ + if scale is None: + scale = k.shape[-1] ** -0.5 + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + + dtype = q.dtype + G = q.shape[2] // k.shape[2] + BS = block_size + k, v, block_indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, block_indices)) + q, k, v = map(lambda x: x.float(), (q, k, v)) + + o = torch.zeros_like(v) + varlen = True + if cu_seqlens is None: + varlen = False + B, T = q.shape[:2] + cu_seqlens = torch.cat([ + block_indices.new_tensor(range(0, B*T, T)), block_indices.new_tensor([B*T]), + ]) + + for i in range(len(cu_seqlens) - 1): + if not varlen: + q_b, k_b, v_b, i_b = q[i], k[i], v[i], block_indices[i] + else: + T = cu_seqlens[i+1] - cu_seqlens[i] + q_b, k_b, v_b, i_b = map(lambda x: x[0][cu_seqlens[i]:cu_seqlens[i+1]], (q, k, v, block_indices)) + + i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS)) + # [T, S*BS, HQ] + i_b = i_b.view(T, block_indices.shape[2], -1).transpose(1, 2) + for i_q in range(T): + # [HQ, D] + q_i = q_b[i_q] * scale + # [S*BS, HQ] + i_i = i_b[i_q] + # [S*BS, HQ, -1] + k_i, v_i = map(lambda x: x.gather(0, i_i.clamp(0, T-1).unsqueeze(-1).expand(*i_i.shape, x.shape[-1])), (k_b, v_b)) + # [S*BS, HQ] + attn = torch.einsum('h d, n h d -> n h', q_i, k_i).masked_fill(i_i > i_q, float('-inf')).softmax(0) + if not varlen: + o[i, i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) + else: + o[0][cu_seqlens[i]+i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) + + return o.to(dtype) diff --git a/fla/ops/nsa/parallel.py b/fla/ops/nsa/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..42af694d86fa9042e1675b284d7a461cfdf81461 --- /dev/null +++ b/fla/ops/nsa/parallel.py @@ -0,0 +1,882 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.attn.parallel import parallel_attn_bwd_preprocess +from fla.ops.nsa.compression import parallel_nsa_compression +from fla.ops.nsa.utils import _bitonic_merge +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices +from fla.ops.utils.op import exp, log +from fla.ops.utils.pooling import mean_pooling +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning, + ) + flash_attn_func = None + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK'], + **autotune_cache_kwargs, +) +@triton.jit +def parallel_nsa_kernel_topk( + q, + k, + lse, + scale, + block_indices, + cu_seqlens, + token_indices, + chunk_offsets, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + S: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + ################################ + # 1. lse computation + ################################ + if lse is not None: + b_lse = tl.load(lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G)) + else: + # max scores for the current block + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + # lse = log(acc) + m + b_acc = tl.zeros([G], dtype=tl.float32) + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BC] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + + b_mp = b_m + if NC == 0: + b_lse = tl.zeros([G], dtype=tl.float32) + else: + b_lse = b_m + log(b_acc) + + ################################ + # 2. topk selection + ################################ + # [BC] + b_i = tl.full([BC], -1, dtype=tl.float32) + o_i = tl.zeros([BC], dtype=tl.int32) + m_i = tl.arange(0, BC) < BC//2 + + IC = i_t // BS + for i_c in range(0, tl.cdiv(i_t + 1, BS), BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where(o_c < IC, b_s, float('-inf')) + # [G, BC] + # the 1st and the last 2 blocks are always selected + b_p = tl.where((o_c == 0) | ((o_c == IC - 1) | (o_c == IC)), 1., exp(b_s - b_lse[:, None])) + # the importance scores of the current block + # [BC] + b_i, b_ip = tl.sum(b_p, 0), b_i + # blocks with index < 0 will be skipped + o_i, o_ip = tl.where(o_c <= IC, o_c, -1), o_i + + n_dims: tl.constexpr = tl.standard._log2(b_i.shape[0]) + for i in tl.static_range(1, n_dims): + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), i, 2, n_dims) + + if i_c != 0: + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, False, n_dims) + b_i_new = b_ip * m_i + b_i * (1 - m_i) + o_i_new = o_ip * m_i + o_i * (1 - m_i) + b_i, o_i = _bitonic_merge(b_i_new, o_i_new.to(tl.int32), n_dims, True, n_dims) + else: + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, True, n_dims) + + m_top = tl.arange(0, BC//S) == 0 + b_top = tl.sum(m_top[:, None] * tl.reshape(o_i, [BC//S, S]), 0) + + p_b = tl.make_block_ptr(block_indices + (bos + i_t) * H*S, (H*S,), (1,), (i_h * S,), (S,), (0,)) + tl.store(p_b, b_top.to(p_b.dtype.element_ty)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit +def parallel_nsa_fwd_kernel( + q, + k, + v, + o, + lse, + scale, + block_indices, + block_counts, + cu_seqlens, + token_indices, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr, +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + block_indices += (bos + i_t) * H*S + i_h * S + + if USE_BLOCK_COUNTS: + NS = tl.load(block_counts + (bos + i_t) * H + i_h) + else: + NS = S + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_o = tl.make_block_ptr(o + (bos + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_o = tl.zeros([G, BV], dtype=tl.float32) + + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + b_acc = tl.zeros([G], dtype=tl.float32) + for i in range(NS): + i_s = tl.load(block_indices + i).to(tl.int32) * BS + if i_s <= i_t and i_s >= 0: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [G, BS] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((i_t >= (i_s + tl.arange(0, BS)))[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BS] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [G, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + b_o = b_o / b_acc[:, None] + b_m += log(b_acc) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), +}) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_kernel_mask( + block_indices, + block_counts, + block_mask, + T, + H: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + NS: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr, +): + i_t, i_b, i_hs = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_h, i_s = i_hs // S, i_hs % S + + b_i = tl.load(block_indices + i_b * T * H * S + i_t * H * S + i_h * S + i_s) + if USE_BLOCK_COUNTS: + b_m = b_i * BS <= i_t and i_s < tl.load(block_counts + i_b * T * H + i_t * H + i_h) + else: + b_m = b_i * BS <= i_t + + if b_i < NS and b_i >= 0: + tl.store(block_mask + i_b * T * H * NS + i_t * H * NS + i_h * NS + b_i, b_m.to(block_mask.dtype.element_ty)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_bwd_kernel_dq( + q, + k, + v, + lse, + delta, + do, + dq, + scale, + block_indices, + block_counts, + cu_seqlens, + token_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr, +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + q += (bos + i_t) * HQ*K + do += (bos + i_t) * HQ*V + lse += (bos + i_t) * HQ + delta += (bos + i_t) * HQ + dq += (i_v * all + bos + i_t) * HQ*K + block_indices += (bos + i_t) * H*S + i_h * S + + if USE_BLOCK_COUNTS: + NS = tl.load(block_counts + (bos + i_t) * H + i_h) + else: + NS = S + + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + + p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + i_h * G + tl.arange(0, G) + p_delta = delta + i_h * G + tl.arange(0, G) + + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + + # [G, BK] + b_dq = tl.zeros([G, BK], dtype=tl.float32) + for i in range(NS): + i_s = tl.load(block_indices + i).to(tl.int32) * BS + if i_s <= i_t and i_s >= 0: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # [G, BS] + b_s = tl.dot(b_q, b_k) + b_p = exp(b_s - b_lse[:, None]) + b_p = tl.where((i_t >= (i_s + tl.arange(0, BS)))[None, :], b_p, 0) + + # [G, BV] @ [BV, BS] -> [G, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [G, BS] @ [BS, BK] -> [G, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq *= scale + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_bwd_kernel_dkv( + q, + k, + v, + lse, + delta, + do, + dk, + dv, + block_mask, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + M: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_s, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + all = B * T + if IS_VARLEN: + i_n, i_s = tl.load(chunk_indices + i_s * 2).to(tl.int32), tl.load(chunk_indices + i_s * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s * BS, 0), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s * BS, i_v * BV), (BS, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_v * all * H + bos * H + i_h) * K, (T, K), (H*K, 1), (i_s * BS, 0), (BS, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s * BS, i_v * BV), (BS, BV), (1, 0)) + + # [BS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BS, BK], dtype=tl.float32) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BS, BV], dtype=tl.float32) + + for i in range(i_s * BS, T): + b_m = tl.load(block_mask + (bos + i) * H*M + i_h * M + i_s) + if b_m: + p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) + p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + # [BS, G] + b_s = tl.dot(b_k, tl.trans(b_q)) + b_p = exp(b_s - b_lse[None, :]) + b_p = tl.where((i >= (i_s * BS + tl.arange(0, BS)))[:, None], b_p, 0) + # [BS, G] @ [G, BV] -> [BS, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BS, BV] @ [BV, G] -> [BS, G] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BS, G] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BS, G] @ [G, BK] -> [BS, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +def parallel_nsa_topk( + q: torch.Tensor, + k: torch.Tensor, + lse: torch.Tensor, + block_counts: torch.LongTensor | int, + block_size: int = 64, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, +) -> torch.LongTensor: + B, T, HQ, K = q.shape + H = k.shape[2] + G = HQ // H + # the number of selected blocks for each token + S = block_counts if isinstance(block_counts, int) else block_counts.max().item() + S = triton.next_power_of_2(S) + # here we set BC = BS, but beware that they can be chosen separately if required + BC = BS = block_size + BK = max(triton.next_power_of_2(K), 16) + assert BC >= 2 * S, f"BC ({BC}) must be greater than or equal to 2 * S ({S})" + + block_indices = torch.zeros(B, T, H, S, dtype=torch.int32, device=q.device) + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) if cu_seqlens is not None else None + grid = (T, B * H) + # the 1st and the last 2 blocks are always selected + parallel_nsa_kernel_topk[grid]( + q=q, + k=k, + lse=lse, + scale=scale, + block_indices=block_indices, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + S=S, + BC=BC, + BS=BS, + BK=BK, + ) + return block_indices + + +def parallel_nsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_indices: torch.LongTensor, + block_counts: torch.LongTensor | int, + block_size: int, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + token_indices: torch.LongTensor | None = None, +): + B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] + HQ = q.shape[2] + G = HQ // H + BS = block_size + if check_shared_mem('hopper', q.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + else: + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, "The key dimension can not be larger than 256" + + grid = (T, NV, B * H) + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + parallel_nsa_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + lse=lse, + scale=scale, + block_indices=block_indices, + block_counts=block_counts, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + S=S, + BS=BS, + BK=BK, + BV=BV, + ) + return o, lse + + +def parallel_nsa_block_mask( + block_indices: torch.LongTensor, + block_counts: torch.LongTensor | int, + cu_seqlens: torch.LongTensor, + block_size: int, +): + B, T, H, S = block_indices.shape + BS = block_size + if cu_seqlens is not None: + NS = triton.cdiv(prepare_lens(cu_seqlens).max().item(), BS) + else: + NS = triton.cdiv(T, BS) + block_mask = torch.zeros(B, T, H, NS, dtype=torch.bool, device=block_indices.device) + + parallel_nsa_kernel_mask[(T, B, H*S)]( + block_indices=block_indices, + block_counts=block_counts, + block_mask=block_mask, + T=T, + H=H, + S=S, + BS=BS, + NS=NS, + ) + return block_mask + + +def parallel_nsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + block_indices: torch.Tensor, + block_counts: torch.LongTensor | int, + block_size: int = 64, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + token_indices: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] + HQ = q.shape[2] + G = HQ // H + BS = block_size + BK = max(triton.next_power_of_2(K), 16) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + NV = triton.cdiv(V, BV) + + delta = parallel_attn_bwd_preprocess(o, do) + + dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) + grid = (T, NV, B * H) + parallel_nsa_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dq=dq, + block_indices=block_indices, + block_counts=block_counts, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + S=S, + BS=BS, + BK=BK, + BV=BV, + ) + dq = dq.sum(0) + + if cu_seqlens is not None: + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BS) + NS = len(chunk_indices) + else: + NS = triton.cdiv(T, BS) + + # [B, T, H, M] + block_mask = parallel_nsa_block_mask(block_indices, block_counts, cu_seqlens, block_size) + dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) + dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) + + grid = (NV, NS, B * H) + parallel_nsa_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + block_mask=block_mask, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + M=block_mask.shape[-1], + BS=BS, + BK=BK, + BV=BV, + ) + dk = dk.sum(0) + return dq, dk, dv + + +@torch.compile +class ParallelNSAFunction(torch.autograd.Function): + + @staticmethod + @contiguous + @autocast_custom_fwd + def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens): + ctx.dtype = q.dtype + + # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence + # for example, if the passed `cu_seqlens` is [0, 2, 6], + # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + token_indices = prepare_token_indices(cu_seqlens) if cu_seqlens is not None else None + + o, lse = parallel_nsa_fwd( + q=q, + k=k, + v=v, + block_indices=block_indices, + block_counts=block_counts, + block_size=block_size, + scale=scale, + cu_seqlens=cu_seqlens, + token_indices=token_indices, + ) + ctx.save_for_backward(q, k, v, o, lse) + ctx.block_indices = block_indices + ctx.block_counts = block_counts + ctx.cu_seqlens = cu_seqlens + ctx.token_indices = token_indices + ctx.block_size = block_size + ctx.scale = scale + return o.to(q.dtype) + + @staticmethod + @contiguous + @autocast_custom_bwd + def backward(ctx, do): + q, k, v, o, lse = ctx.saved_tensors + dq, dk, dv = parallel_nsa_bwd( + q=q, + k=k, + v=v, + o=o, + lse=lse, + do=do, + block_indices=ctx.block_indices, + block_counts=ctx.block_counts, + block_size=ctx.block_size, + scale=ctx.scale, + cu_seqlens=ctx.cu_seqlens, + token_indices=ctx.token_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), None, None, None, None, None, None, None, None + + +def parallel_nsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cmp: torch.Tensor | None = None, + g_slc: torch.Tensor | None = None, + g_swa: torch.Tensor | None = None, + block_indices: torch.LongTensor | None = None, + block_counts: torch.LongTensor | int = 16, + block_size: int = 64, + window_size: int = 0, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g_cmp (torch.Tensor): + Gate score for compressed attention of shape `[B, T, HQ]`. + g_slc (torch.Tensor): + Gate score for selected attention of shape `[B, T, HQ]`. + g_swa (torch.Tensor): + Gate score for sliding attentionof shape `[B, T, HQ]`. + block_indices (torch.LongTensor): + Block indices of shape `[B, T, H, S]`. + `S` is the number of selected blocks for each query token, which is set to 16 in the paper. + If `g_cmp` is provided, the passed `block_indices` will be ignored. + block_counts (Optional[Union[torch.LongTensor, int]]): + Number of selected blocks for each query. + If a tensor is provided, with shape `[B, T, H]`, + each query can select the same number of blocks. + If not provided, it will default to 16. + block_size (int): + Selected block size. Default: 64. + window_size (int): + Sliding window size. Default: 0. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + """ + assert block_counts is not None, "block counts must be provided for selection" + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + assert q.shape[2] % (k.shape[2] * 16) == 0, "Group size must be a multiple of 16 in NSA" + + k_cmp, v_cmp = mean_pooling(k, block_size, cu_seqlens), mean_pooling(v, block_size, cu_seqlens) + o_cmp, lse_cmp = None, None + if g_cmp is not None: + o_cmp, lse_cmp = parallel_nsa_compression( + q=q, + k=k_cmp, + v=v_cmp, + block_size=block_size, + scale=scale, + cu_seqlens=cu_seqlens, + ) + if block_indices is not None: + warnings.warn("`block_indices` will be ignored when `g_cmp` is provided") + block_indices = parallel_nsa_topk( + q=q, + k=k_cmp, + lse=lse_cmp, + block_counts=block_counts, + block_size=block_size, + scale=scale, + cu_seqlens=cu_seqlens, + ) + o = o_slc = ParallelNSAFunction.apply(q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens) + if g_slc is not None: + o = o_slc * g_slc.unsqueeze(-1) + if o_cmp is not None: + o = torch.addcmul(o, o_cmp, g_cmp.unsqueeze(-1)) + if window_size > 0: + if cu_seqlens is not None: + max_seqlen = q.shape[1] + o_swa = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(window_size-1, 0), + ).unsqueeze(0) + else: + o_swa = flash_attn_func( + q, k, v, + causal=True, + window_size=(window_size-1, 0), + ) + o = torch.addcmul(o, o_swa, g_swa.unsqueeze(-1)) + return o diff --git a/fla/ops/nsa/utils.py b/fla/ops/nsa/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0eadb7798292cd92a2220366a2e63c2a2f7e59f7 --- /dev/null +++ b/fla/ops/nsa/utils.py @@ -0,0 +1,91 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Implements argsort based on bitonic sort. +# [What is bitonic sort?](https://en.wikipedia.org/wiki/Bitonic_sorter) + +# Code adapted from https://github.com/triton-lang/triton/issues/3698#issuecomment-2067681396 + + +import triton +import triton.language as tl + +from fla.ops.utils.op import log2 + + +@triton.jit +def _compare_and_swap( + x, + ids, + flip, + i: tl.constexpr, + n_dims: tl.constexpr, +): + n_outer: tl.constexpr = x.numel >> n_dims + shape: tl.constexpr = [n_outer * 2**i, 2, 2**(n_dims - i - 1)] + y = tl.reshape(x, shape) + # slice left/right with 'stride' 2**(n_dims - i - 1) + mask = tl.arange(0, 2)[None, :, None] + left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype) + right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype) + left = tl.reshape(left, x.shape) + right = tl.reshape(right, x.shape) + # idx + y_idx = tl.reshape(ids, shape) + left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape) + right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape) + left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype) + right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype) + # actual compare-and-swap + idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) + ileft = left.to(idtype, bitcast=True) + iright = right.to(idtype, bitcast=True) + ix = x.to(idtype, bitcast=True) + + cond = (left > right) != flip + ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix)) + new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids)) + return ret.to(x.dtype, bitcast=True), new_ids + + +@triton.jit +def _bitonic_merge( + x, + ids, + stage: tl.constexpr, + order: tl.constexpr, + n_dims: tl.constexpr, +): + n_outer: tl.constexpr = x.numel >> n_dims + tl.static_assert(stage <= n_dims) + # flip denotes whether to re-arrange sub-sequences of elements in ascending or + # descending order. + # if flip = 00000000... then all elements will be re-arranged ascendingly at this stage + # if flip = 00110011... then all the elements will be re-arranged alternatingly (with + # a stride of 2) at this stage + if order == 2: + shape: tl.constexpr = [n_outer * 2**(n_dims - 1 - stage), 2, 2**stage] + flip = tl.reshape(tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape) + else: + flip = order + # perform `stage` rounds of `compare-and-swap` + for i in tl.static_range(stage): + x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims) + return x, ids + + +@triton.jit +def argsort( + x, + ids, + dim: tl.constexpr = None, + descending: tl.constexpr = tl.core.CONSTEXPR_0, +): + # handle default dimension or check that it is the most minor dim + _dim: tl.constexpr = len(x.shape) - 1 if dim is None else dim + tl.static_assert(_dim == len(x.shape) - 1, "only minor dimension is currently supported") + # iteratively run bitonic merge-sort steps + n_dims: tl.constexpr = log2(x.shape[_dim]) + + for i in tl.static_range(1, n_dims + 1): + x, ids = _bitonic_merge(x, ids, i, 2 if i < n_dims else descending, n_dims) + return x, ids diff --git a/fla/ops/path_attn/__init__.py b/fla/ops/path_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..854267d78cc92df6da1d3bfd1d57a51c7f54d364 --- /dev/null +++ b/fla/ops/path_attn/__init__.py @@ -0,0 +1,8 @@ + +from .naive import naive_path_attn +from .parallel import parallel_path_attn + +__all__ = [ + 'naive_path_attn', + 'parallel_path_attn', +] diff --git a/fla/ops/path_attn/cumprod_householder_bwd.py b/fla/ops/path_attn/cumprod_householder_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..d13ba7cdde2f000ee6e71397cae538aca5ee6750 --- /dev/null +++ b/fla/ops/path_attn/cumprod_householder_bwd.py @@ -0,0 +1,142 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.utils import check_shared_mem + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_cumprod_householder_bwd_kernel( + hc_suffix, dhc_whole, + k, dk, w1, w2, dw1, dw2, dk_new, + cu_seqlens, split_indices, chunk_offsets, split_offsets, + BT: tl.constexpr, # previous small chunk size + K: tl.constexpr, + BK: tl.constexpr, + T: tl.constexpr, + S: tl.constexpr, + G: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_ss, i_hq = tl.program_id(0), tl.program_id(1) + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh_large = tl.load(split_offsets + i_n).to(tl.int32) + else: + NS = tl.cdiv(T, S) + i_n, i_s = i_ss // NS, i_ss % NS + bos, eos = i_n * T, i_n * T + T + boh = i_n * tl.cdiv(T, BT) + boh_large = i_n * tl.cdiv(T, S) + + # offset calculations + dhc_whole += ((boh_large + i_s) * HQ + i_hq) * K * K + hc_suffix += ((boh + tl.cdiv(i_s * S, BT)) * H + i_h) * K * K + k += (bos * H + i_h) * K + w1 += (bos * H + i_h) * K + w2 += (bos * H + i_h) * K + dw1 += (bos * HQ + i_hq) * K + dw2 += (bos * HQ + i_hq) * K + + # dh += ((boh + tl.cdiv(i_s * S, BT)) * HQ + i_hq) * K * K + dk += (bos * HQ + i_hq) * K + dk_new += (bos * HQ + i_hq) * K + + stride_h = H * K * K + NT_small = tl.cdiv(min(S, T-i_s*S), BT) + p_dhc_whole = tl.make_block_ptr(dhc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + b_dhc = tl.zeros([BK, BK], dtype=tl.float32) + b_dhc += tl.load(p_dhc_whole, boundary_check=(0, 1)) + + # calculate dh + for i_t_small in range(0, NT_small): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_dk_new = tl.make_block_ptr(dk_new, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_hc = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)) + + p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_hc = tl.load(p_hc, boundary_check=(0, 1)) + + b_dk_new = b_dk - tl.dot(b_dk.to(b_hc.dtype), b_hc) + tl.store(p_dk_new, b_dk_new.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) + + b_dh = b_dhc - tl.dot(tl.trans(b_hc), b_dhc.to(b_hc.dtype)) + b_dw2 = tl.dot(b_w1, b_dh.to(b_w1.dtype)) + b_dw1 = tl.dot(b_w2, tl.trans(b_dh.to(b_w2.dtype))) + + p_dw1 = tl.make_block_ptr(dw1, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_dw2 = tl.make_block_ptr(dw2, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + + tl.store(p_dw1, b_dw1.to(dw1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw2, b_dw2.to(dw2.dtype.element_ty), boundary_check=(0, 1)) + + b_dhc = b_dhc - tl.dot(tl.dot(b_dhc.to(b_w2.dtype), tl.trans(b_w2)).to(b_w1.dtype), b_w1) + b_dhc -= tl.dot(tl.trans(b_dk).to(b_k.dtype), b_k) + + +def chunk_cumprod_householder_bwd_fn( + w1: torch.Tensor, + w2: torch.Tensor, + hc_suffix: torch.Tensor, + dhc_whole: torch.Tensor, + k: torch.Tensor, + dk: torch.Tensor, + S: int, # split size, aka large chunk size + BT: int, # small chunk size + cu_seqlens: torch.Tensor = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, HQ, K = dk.shape + H = k.shape[2] + G = HQ // H + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, S) + split_indices = chunk_indices + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None + split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None + + if cu_seqlens is None: + N = B + NS = N * triton.cdiv(T, S) + else: + N = len(cu_seqlens) - 1 + NS = split_offsets[-1].item() + + grid = (NS, HQ) + dw1 = torch.empty_like(dk, dtype=torch.float32) + dw2 = torch.empty_like(dk, dtype=torch.float32) + dk_new = torch.empty_like(dk, dtype=torch.float32) + + chunk_cumprod_householder_bwd_kernel[grid]( + hc_suffix=hc_suffix, dhc_whole=dhc_whole, + k=k, dk=dk, w1=w1, w2=w2, dw1=dw1, dw2=dw2, dk_new=dk_new, + cu_seqlens=cu_seqlens, + split_indices=split_indices, chunk_offsets=chunk_offsets, split_offsets=split_offsets, + BT=BT, K=K, G=G, H=H, HQ=HQ, BK=K, + T=T, S=S, + # SY (2025/07/08): I don't know why when K == 128 if I set num_warps=4 the result would be completely wrong + num_warps=8 if K == 128 else 4, + num_stages=2 if check_shared_mem('ampere') else 1, + ) + return dw1, dw2, dk_new diff --git a/fla/ops/path_attn/cumprod_householder_fwd.py b/fla/ops/path_attn/cumprod_householder_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..5bcb2413e8b60b0ed97ac752d62ab443601a9d5e --- /dev/null +++ b/fla/ops/path_attn/cumprod_householder_fwd.py @@ -0,0 +1,122 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.utils import check_shared_mem + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit +def chunk_cumprod_householder_fwd_kernel( + k, + k_new, + w1, + w2, + hc_suffix, + hc_whole, + cu_seqlens, + split_indices, + chunk_offsets, + split_offsets, + BT: tl.constexpr, # small chunk size + K: tl.constexpr, + H: tl.constexpr, + BK: tl.constexpr, + T: tl.constexpr, + S: tl.constexpr, # split size, aka large chunk size + IS_VARLEN: tl.constexpr, +): + i_ss, i_h = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NS = tl.cdiv(T, S) + + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh_large = tl.load(split_offsets + i_n).to(tl.int32) + else: + NS = tl.cdiv(T, S) + i_n, i_s = i_ss // NS, i_ss % NS + bos, eos = i_n * T, i_n * T + T + + boh = i_n * tl.cdiv(T, BT) + boh_large = i_n * tl.cdiv(T, S) + + NT_small = tl.cdiv(min(S, T-i_s*S), BT) + stride_h = H*K*K + + # offset calculations + hc_whole += ((boh_large + i_s) * H + i_h) * K * K + hc_suffix += ((boh + tl.cdiv(i_s * S, BT)) * H + i_h) * K * K + + k += (bos * H + i_h) * K + k_new += (bos * H + i_h) * K + w1 += (bos * H + i_h) * K + w2 += (bos * H + i_h) * K + + b_h = tl.zeros([BK, BK], dtype=tl.float32) + for i_t_small in range(NT_small-1, -1, -1): + p_hc_suffix = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + tl.store(p_hc_suffix, b_h.to(hc_suffix.dtype.element_ty), boundary_check=(0, 1)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = (b_k - tl.dot(b_k, tl.trans(b_h.to(b_k.dtype)))) + p_w1 = tl.make_block_ptr(w1, (K, T), (1, H*K), (0, i_s * S + i_t_small * BT), (BK, BT), (0, 1)) + p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_v_new = (b_w1 - tl.dot(b_h.to(b_w1.dtype), b_w1)).to(b_w2.dtype) + b_h += tl.dot(b_v_new, b_w2) + p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) + tl.store(p_k_new, b_k.to(k_new.dtype.element_ty), boundary_check=(0, 1)) + + p_hc_whole = tl.make_block_ptr(hc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + tl.store(p_hc_whole, b_h.to(hc_whole.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_cumprod_householder_fwd_fn( + k: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + S: int, # split size, aka large chunk size + BT: int, # small chunk size + cu_seqlens: torch.Tensor = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = k.shape + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, S) + split_indices = chunk_indices + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None + split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None + + if cu_seqlens is None: + N = B + NS = N * triton.cdiv(T, S) + NT = N * triton.cdiv(T, BT) + else: + N = len(cu_seqlens) - 1 + NS = split_offsets[-1] + NT = chunk_offsets[-1] + + grid = (NS, H) + hc_whole = torch.empty((NS, H, K, K), device=k.device, dtype=w1.dtype) + k_new = torch.empty_like(k, dtype=k.dtype) + hc_suffix = torch.empty((NT, H, K, K), device=k.device, dtype=w1.dtype) + chunk_cumprod_householder_fwd_kernel[grid]( + k=k, k_new=k_new, w1=w1, w2=w2, hc_whole=hc_whole, hc_suffix=hc_suffix, + cu_seqlens=cu_seqlens, + split_indices=split_indices, chunk_offsets=chunk_offsets, split_offsets=split_offsets, + BT=BT, K=K, H=H, BK=K, + T=T, S=S, + # SY (2025/07/08): I don't know why when K == 128 if I set num_warps=4 the result would be completely wrong + num_warps=8 if K == 128 else 4, + num_stages=3 if check_shared_mem('ampere') else 1, + ) + return k_new, hc_suffix, hc_whole diff --git a/fla/ops/path_attn/intra_chunk_preprocess_bwd.py b/fla/ops/path_attn/intra_chunk_preprocess_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..9f988bd4cad88b349c22d3d6b11ea994e5c7e922 --- /dev/null +++ b/fla/ops/path_attn/intra_chunk_preprocess_bwd.py @@ -0,0 +1,144 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import check_shared_mem + + +# episold +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['offsets'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def intra_chunk_preprocess_bwd_kernel( + q, k, w, w2, beta, + AT, + dA_local, dq, dq_new, dk, dk_new, dw, dbeta, dw1, dw2, T, + offsets, indices, + HQ: tl.constexpr, G: tl.constexpr, H: tl.constexpr, + K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw_beta = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_dT = tl.zeros([BT, BT], dtype=tl.float32) + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + (bos * H + i_h), (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) + p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_Twb = tl.load(p_w2, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_T = tl.load(p_T, boundary_check=(0, 1)) + b_w_beta = (b_w * b_beta[:, None]).to(b_w.dtype) + + o_i = tl.arange(0, BT) + b_qw = tl.where(o_i[:, None] >= o_i[None, :], tl.dot(b_q, tl.trans(b_w)), 0).to(b_q.dtype) + b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w_beta, tl.trans(b_k)), 0).to(b_k.dtype) + b_Twbk = tl.dot(b_T, b_wbk).to(b_w.dtype) + + p_dA_local = tl.make_block_ptr(dA_local + (bos * HQ + i_hq) * BT, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_dA_local = tl.load(p_dA_local, boundary_check=(0, 1)) + + # # Twb part qw part. + p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dq = tl.load(p_dq, boundary_check=(0, 1)) + + p_dw1 = tl.make_block_ptr(dw1 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dw += tl.load(p_dw1, boundary_check=(0, 1)) + + b_dqw = -tl.dot(b_dA_local, tl.trans(b_Twbk)) - tl.dot(b_dq.to(b_Twb.dtype), tl.trans(b_Twb)) + p_dw2 = tl.make_block_ptr(dw2 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dTwb = -tl.dot(tl.trans(b_qw), b_dq) + tl.load(p_dw2, boundary_check=(0, 1)) + b_dT += tl.dot(b_dTwb.to(b_w_beta.dtype), tl.trans(b_w_beta)) + b_dw_beta += tl.dot(tl.trans(b_T), b_dTwb.to(b_T.dtype)) + + b_dqw = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dqw, 0) + b_dq += tl.dot(b_dA_local.to(b_k.dtype), b_k) + b_dq += tl.dot(b_dqw.to(b_w.dtype), b_w) + b_dw += tl.dot(tl.trans(b_dqw.to(b_q.dtype)), b_q) + p_q_new = tl.make_block_ptr(dq_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) + + # Twbk part + p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dk = tl.load(p_dk, boundary_check=(0, 1)) + b_dTwbk = -tl.dot(tl.trans(b_qw), b_dA_local.to(b_qw.dtype)) - tl.dot(b_w, tl.trans(b_dk.to(b_w.dtype))) + b_dw -= tl.dot(b_Twbk, b_dk.to(b_w.dtype)) + b_dT += tl.dot(b_dTwbk.to(b_wbk.dtype), tl.trans(b_wbk)) + b_dwbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(tl.trans(b_T), b_dTwbk.to(b_T.dtype)), 0).to(b_w.dtype) + b_dw_beta += tl.dot(b_dwbk, b_k) + + b_dk += tl.dot(tl.trans(b_dwbk), b_w_beta) + b_dk += tl.dot(tl.trans(b_dA_local), b_q) + p_dk_new = tl.make_block_ptr(dk_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dk_new, b_dk.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) + + # matrix inverse's gradient + p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (BT, T), (1, BT*H), (0, i_t * BT), (BT, BT), (0, 1)) + b_Tt = tl.load(p_T, boundary_check=(0, 1)) + b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dT, 0).to(b_w.dtype) + b_dT = tl.dot(b_Tt, b_dT).to(b_w.dtype) + b_dT = tl.dot(b_dT, b_Tt) + b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dT, 0).to(b_k.dtype) + + b_dw_beta += tl.dot(b_dT, b_w) + b_dw += tl.dot(tl.trans(b_dT), b_w_beta) + b_dw += b_dw_beta * b_beta[:, None] + b_dbeta = tl.sum(b_dw_beta * b_w, axis=1) + + p_dw = tl.make_block_ptr(dw + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dw, b_dw.to(dw.dtype.element_ty), boundary_check=(0, 1)) + p_dbeta = tl.make_block_ptr(dbeta + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + tl.store(p_dbeta, b_dbeta.to(dbeta.dtype.element_ty), boundary_check=(0, )) + + +def intra_chunk_preprocess_bwd_fn(q, k, w, w2, beta, + dq, dk, dA_local, + dw1, dw2, + A, L, D, do, scale, cu_seqlens=None, + chunk_indices: torch.LongTensor | None = None): + BT = A.shape[-1] + HQ = q.shape[-2] + B, T, H, K = k.shape + G = HQ//H + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + grid = (NT, B*HQ) + # better precision because h would be of norm smaller than 1 anyways + + dbeta = torch.empty(B, T, HQ, device=q.device, dtype=k.dtype if G == 1 else torch.float32) + dw = torch.empty(B, T, HQ, K, device=q.device, dtype=k.dtype if G == 1 else torch.float32) + dk_new = torch.empty_like(dk, dtype=k.dtype if G == 1 else torch.float32) # float32 reduction + dq_new = torch.empty_like(dq, dtype=q.dtype) + + intra_chunk_preprocess_bwd_kernel[grid]( + q=q, k=k, w=w, w2=w2, beta=beta, + AT=A, + dA_local=dA_local, dq=dq, dq_new=dq_new, dk=dk, dk_new=dk_new, dw=dw, dbeta=dbeta, dw1=dw1, dw2=dw2, T=T, + offsets=cu_seqlens, indices=indices, + HQ=HQ, G=G, H=H, + K=K, BT=BT, BK=triton.next_power_of_2(K), + num_stages=3 if check_shared_mem('hopper') else 1, + ) + return dq_new, dk_new, dbeta, dw diff --git a/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py b/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..acdce96bb7f735b29cb3400239de2659e26ab636 --- /dev/null +++ b/fla/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py @@ -0,0 +1,200 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets + + +@triton.heuristics({ + "USE_GATE": lambda args: args['g_cumsum'] is not None, + "IS_VARLEN": lambda args: args['offsets'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_transform_qk_bwd_kernel_prepare( + q, + k, + v, + w, + beta, + g_cumsum, + L, + D, + h, + q_new, + k_new, + AT, + dA_local, + dv, + do, + dg_cumsum, + scale, + indices, # varlen helper + offsets, # varlen helper + chunk_offsets, # varlen helper + T, + G: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATE: tl.constexpr, + RETURN_H: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + sm_scale = scale * 1.44269504 + # offset calculations + dA_local += (bos*HQ + i_hq) * BT + AT += (bos*H + i_h) * BT + q += (bos*HQ + i_hq) * K + q_new += (bos*HQ + i_hq) * K + k += (bos*H + i_h) * K + k_new += (bos*H + i_h) * K + w += (bos*H + i_h) * K + v += (bos*H + i_h) * V + do += (bos*HQ + i_hq) * V + dv += (bos*HQ + i_hq) * V + beta += (bos*H + i_h) + if RETURN_H: + h += ((boh + i_t) * H + i_h) * K * K + else: + h += (bos*H + i_h) * K + if USE_GATE: + g_cumsum += (bos*HQ + i_hq) + dg_cumsum += (bos*HQ + i_hq) + L += (bos*HQ + i_hq) + D += (bos*HQ + i_hq) + + p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kt = tl.load(p_k, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0, )) + p_T = tl.make_block_ptr(AT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_T = tl.load(p_T, boundary_check=(0, 1)) * b_beta[None, :] + + o_i = tl.arange(0, BT) + m_t = o_i[:, None] >= o_i[None, :] + b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) + b_qwT = tl.dot(b_qw, b_T.to(b_q.dtype)).to(b_q.dtype) + b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w.to(b_kt.dtype), b_kt), 0).to(b_q.dtype) + b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT, b_wbk), 0) + + b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_qwT.dtype)) + p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) + tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + + if i_hq % G == 0: + b_Twb = tl.dot(b_T, b_w) # tf32 + p_h = tl.make_block_ptr(h, (T, K), (K * H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_h, b_Twb.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + b_T_wbk = tl.dot(b_T.to(b_wbk.dtype), b_wbk).to(b_kt.dtype) + p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) + tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) + ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + + if USE_GATE: + p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) + b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) + b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan + + p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_l = tl.load(p_l, boundary_check=(0, )) + p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + delta = tl.load(p_delta, boundary_check=(0, )) + + b_A_softmax = tl.exp2(tl.where(o_i[:, None] >= o_i[None, :], b_A * sm_scale - b_l[:, None], float("-inf"))) + p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv = tl.dot(tl.trans(b_A_softmax.to(b_do.dtype)), b_do) + p_dv = tl.make_block_ptr(dv, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (0, i_t * BT), (BV, BT), (0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dp = tl.dot(b_do, b_v) + b_dA = ((b_dp - delta[:, None]) * b_A_softmax * scale) + if USE_GATE: + b_dgq = tl.sum(b_dA, axis=1) - tl.sum(b_dA, axis=0) + p_dg = tl.make_block_ptr(dg_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + tl.store(p_dg, b_dgq.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dA = tl.make_block_ptr(dA_local, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +def intra_chunk_preprocess_bwd_prepare_fn(q, k, v, w, beta, g_cumsum, A, L, D, do, scale, return_h=True, cu_seqlens=None, + chunk_indices: torch.LongTensor | None = None): + BT = A.shape[-1] + HQ = q.shape[-2] + B, T, H, K = k.shape + G = HQ//H + + V = v.shape[-1] + q_new = torch.empty_like(q) + k_new = torch.empty_like(k) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + grid = (NT, B*HQ) + h = torch.empty_like(w) + dA_local = torch.empty(B, T, HQ, BT, dtype=q.dtype, device=q.device) + dv = torch.empty(B, T, HQ, V, device=q.device, dtype=torch.float32) + dg_cumsum = torch.empty_like(g_cumsum) if g_cumsum is not None else None + + chunk_transform_qk_bwd_kernel_prepare[grid]( + q=q, + k=k, + v=v, + w=w, + beta=beta, + g_cumsum=g_cumsum, + AT=A, + dA_local=dA_local, + dv=dv, + dg_cumsum=dg_cumsum, + do=do, + L=L, + D=D, + h=h, + q_new=q_new, + k_new=k_new, + scale=scale, + offsets=cu_seqlens, + indices=indices, + chunk_offsets=chunk_offsets, + T=T, + H=H, + G=G, + HQ=HQ, + K=K, + V=V, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + BT=BT, + RETURN_H=return_h, + ) + return q_new, k_new, h, dA_local, dv, dg_cumsum diff --git a/fla/ops/path_attn/intra_chunk_preprocess_fwd.py b/fla/ops/path_attn/intra_chunk_preprocess_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..88a10402253c6675e45a1f8cd6be9cce22adf756 --- /dev/null +++ b/fla/ops/path_attn/intra_chunk_preprocess_fwd.py @@ -0,0 +1,173 @@ + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices + + +@triton.heuristics({ + "USE_G": lambda args: args['g_cumsum'] is not None, + "IS_VARLEN": lambda args: args['offsets'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def intra_chunk_preprocess_fwd_kernel( + q, + k, + v, + w, + beta, + g_cumsum, + o, + A, + L, + M, + w2, + q_new, + k_new, + scale, + indices, # varlen helper + offsets, # varlen helper + T, + H: tl.constexpr, + G: tl.constexpr, + HQ: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + sm_scale = scale * 1.44269504 + # offset calculations + A += (bos*H + i_h) * BT + q += (bos*HQ + i_hq) * K + q_new += (bos*HQ + i_hq) * K + k += (bos*H + i_h) * K + k_new += (bos*H + i_h) * K + w2 += (bos*H + i_h) * K + w += (bos*H + i_h) * K + v += (bos*H + i_h) * V + o += (bos*HQ + i_hq) * V + beta += (bos*H + i_h) + if USE_G: + g_cumsum += (bos*HQ + i_hq) + L += (bos*HQ + i_hq) + M += (bos*HQ + i_hq) + + p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) + p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) + p_T = tl.make_block_ptr(A, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kt = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_T = tl.load(p_T, boundary_check=(0, 1)) + b_T = b_T * b_beta[None, :] + + o_i = tl.arange(0, BT) + m_t = o_i[:, None] >= o_i[None, :] + + b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) + b_qwT = tl.dot(b_qw, b_T.to(b_q.dtype)).to(b_q.dtype) + b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w.to(b_q.dtype), b_kt), 0).to(b_q.dtype) + b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT.to(b_q.dtype), b_wbk), 0) + + b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_q.dtype)) + p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) + tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + + if i_hq % G == 0: + b_Twb = tl.dot(b_T, b_w) + p_w2 = tl.make_block_ptr(w2, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_w2, b_Twb.to(p_w2.dtype.element_ty), boundary_check=(0, 1)) + b_T_wbk = tl.dot(b_T.to(b_kt.dtype), b_wbk).to(b_kt.dtype) + p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) + tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) + ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) + b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) + b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan + + b_qkT_softmax = tl.where(o_i[:, None] >= o_i[None, :], b_A * sm_scale, float("-inf")) + m_i = tl.max(b_qkT_softmax, 1) + b_qkT_softmax = tl.math.exp2(b_qkT_softmax - m_i[:, None]) + l_i = tl.sum(b_qkT_softmax, 1) + b_o = tl.dot(b_qkT_softmax.to(b_v.dtype), b_v) + p_o = tl.make_block_ptr(o, (T, V), (V*HQ, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + p_m = tl.make_block_ptr(M, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + tl.store(p_m, m_i.to(p_m.dtype.element_ty), boundary_check=(0,)) + tl.store(p_l, l_i.to(p_l.dtype.element_ty), boundary_check=(0,)) + + +def intra_chunk_preprocess_fwd_fn(q, k, v, w, beta, g_cumsum, A, scale, BT, cu_seqlens, + chunk_indices: torch.LongTensor | None = None): + HQ = q.shape[-2] + B, T, H, K = k.shape + V = v.shape[-1] + q_new = torch.empty_like(q, dtype=torch.float32) # for stability + k_new = torch.empty_like(k) + o = torch.empty(B, T, HQ, V, device=q.device, dtype=torch.float32) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + grid = (NT, B*HQ) + L = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) + M = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) + w2 = torch.empty_like(w) + G = HQ//H + + intra_chunk_preprocess_fwd_kernel[grid]( + q=q, + k=k, + v=v, + w=w, + beta=beta, + g_cumsum=g_cumsum, + o=o, + A=A, + L=L, + M=M, + w2=w2, + q_new=q_new, + k_new=k_new, + scale=scale, + offsets=cu_seqlens, + indices=indices, + T=T, + H=H, + G=G, + HQ=HQ, + K=K, + V=V, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + BT=BT, + num_warps=4 if BT == 64 else 2, + ) + return q_new, k_new, w2, o, L, M diff --git a/fla/ops/path_attn/naive.py b/fla/ops/path_attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..2b73ec5d9093313093528806fadb0487a14a947e --- /dev/null +++ b/fla/ops/path_attn/naive.py @@ -0,0 +1,92 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_path_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float, + chunk_size: int = 64, +): + """ + Reference PyTorch implementation of path attention. + + Args: + q: [B, T, HQ, D] + k: [B, T, H, D] + v: [B, T, H, D] + w: [B, T, H, D] + beta: [B, T, H] + g: [B, T, HQ] + scale: float + chunk_size: int, default 64 + + Returns: + output: [B, T, HQ, D] + """ + original_dtype = q.dtype + HQ = q.shape[2] + H = k.shape[2] + BT = chunk_size + + q, k, v, w, beta, g = map(lambda x: x.to(torch.float).transpose(1, 2), [q, k, v, w, beta, g]) + g_cumsum = g.cumsum(-1) + + q = q.unsqueeze(2).expand(-1, -1, HQ // HQ, -1, -1).flatten(1, 2) + k = k.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) + v = v.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) + w = w.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) + beta = beta.unsqueeze(2).expand(-1, -1, HQ // H, -1).flatten(1, 2) + + b, h, l, _ = q.shape + if l % BT != 0: + padding_size = BT - l % BT + q, k, w = map(lambda x: F.pad(x, (0, 0, 0, padding_size)), [q, k, w]) + beta = F.pad(beta, (0, padding_size)) + + seq_len = q.shape[2] + w_beta = w * beta[..., None] + q, k, w, w_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, w, w_beta]) + + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + T_mat = -(w_beta @ w.transpose(-1, -2)).masked_fill(mask, 0) + + for i in range(1, BT): + T_mat[..., i, :i] = T_mat[..., i, :i].clone() + (T_mat[..., i, :, None].clone() * T_mat[..., :, :i].clone()).sum(-2) + + T_mat = T_mat + torch.eye(BT, dtype=q.dtype, device=q.device) + Twbk = T_mat @ (w_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + qw = (q @ w.transpose(-1, -2)).tril() + Twb = T_mat @ w_beta + A_local = (q @ k.transpose(-1, -2)).tril() - qw @ Twbk + q = q - qw @ Twb + k = k - Twbk.transpose(-1, -2) @ w + H_mat = w.transpose(-1, -2) @ Twb + + A = torch.zeros(b, h, seq_len, seq_len, device=q.device) + q, k, w, w_beta = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, w, w_beta]) + + for i in range(0, seq_len, BT): + q_i = q[:, :, i:i+BT].clone() + for j in range(i - BT, -BT, -BT): + k_j = k[:, :, j:j+BT] + A_ij = q_i @ k_j.transpose(-1, -2) + A[:, :, i:i+BT, j:j+BT] = A_ij + q_i = q_i - q_i @ H_mat[:, :, j // BT] + + for i in range(0, seq_len // BT): + A[:, :, i*BT:i*BT+BT, i*BT:i*BT+BT] = A_local[:, :, i] + + A = A.masked_fill_(~torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)), float("-inf")) + A = A[:, :, :l, :l] + A = A + g_cumsum[..., None] - g_cumsum[..., None, :] + ref_o = (A * scale).softmax(-1).to(v) @ v + + return ref_o.to(original_dtype).transpose(1, 2) diff --git a/fla/ops/path_attn/parallel.py b/fla/ops/path_attn/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..6121eceb2512c554b650574d2c44234b296883a7 --- /dev/null +++ b/fla/ops/path_attn/parallel.py @@ -0,0 +1,277 @@ +# Copyright (c) 2024, Songlin Yang, Yu Zhang + + +import torch +from einops import reduce + +from fla.ops.attn.parallel import parallel_attn_bwd_preprocess +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.path_attn.cumprod_householder_bwd import chunk_cumprod_householder_bwd_fn +from fla.ops.path_attn.cumprod_householder_fwd import chunk_cumprod_householder_fwd_fn +from fla.ops.path_attn.intra_chunk_preprocess_bwd import intra_chunk_preprocess_bwd_fn +from fla.ops.path_attn.intra_chunk_preprocess_bwd_prepare import intra_chunk_preprocess_bwd_prepare_fn +from fla.ops.path_attn.intra_chunk_preprocess_fwd import intra_chunk_preprocess_fwd_fn +from fla.ops.path_attn.parallel_path_bwd_inter_dkv import parallel_path_bwd_dkv_fn +from fla.ops.path_attn.parallel_path_bwd_inter_dqh import parallel_path_bwd_dq_fn +from fla.ops.path_attn.parallel_path_bwd_intra import parallel_path_bwd_intra_chunk_fn +from fla.ops.path_attn.parallel_path_fwd import parallel_path_fwd_fn +from fla.ops.path_attn.prepare_k_cache import prepare_k_cache_fn +from fla.ops.path_attn.transform_q import transform_q_fwd_fn +from fla.ops.utils.cumsum import chunk_global_cumsum +from fla.ops.utils.solve_tril import solve_tril +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, input_guard + + +class ParallelPATHAttentionFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, w, beta, g, scale, cu_seqlens, use_cache=False): + + g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, output_dtype=torch.float32) if g is not None else None + BS = 64 if check_shared_mem('hopper') else 32 + BT = 128 if check_shared_mem('ampere') else 64 + + A = chunk_scaled_dot_kkt_fwd( + k=w, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_size=BS, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + output_dtype=w.dtype, # force fp32? + ) + q_new, k_new, w2, o, L, M = intra_chunk_preprocess_fwd_fn( + q=q, + k=k, + v=v, + w=w, + beta=beta, + g_cumsum=g_cumsum, + A=A, + scale=scale, + BT=BS, + cu_seqlens=cu_seqlens, + ) + w_fp16 = w.to(torch.float16) + w2_fp16 = w2.to(torch.float16) + o, L = parallel_path_fwd_fn( + q=q_new, + k=k_new, + v=v, + L=L, + w1=w_fp16, + w2=w2_fp16, + M=M, + o=o, + g_cumsum=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + BT=BT, + BS=BS, + ) + k_cache = prepare_k_cache_fn(k=k_new, w1=w, w2=w2, cu_seqlens=cu_seqlens, BS=BS, use_cache=use_cache) + ctx.save_for_backward(q, k, v, w, g_cumsum, o, beta, L, A) + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o, k_cache + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dk_new): + q, k, v, w, g_cumsum, o, beta, L, A = ctx.saved_tensors + BT = 128 if check_shared_mem('ampere') else 64 + BS = 64 if check_shared_mem('hopper') else 32 + S = 512 + cu_seqlens = ctx.cu_seqlens + delta = parallel_attn_bwd_preprocess(o, do) + + q_new, k_new, h, dA_local, dv, dg_cumsum = intra_chunk_preprocess_bwd_prepare_fn( + q=q, + k=k, + v=v, + w=w, + beta=beta, + g_cumsum=g_cumsum, + A=A, + L=L, + D=delta, + do=do, + scale=ctx.scale, + cu_seqlens=cu_seqlens, + return_h=False, + ) + w_fp16 = w.to(torch.float16) + h_fp16 = h.to(torch.float16) + k_new_large, hc_suffix, hc_whole = chunk_cumprod_householder_fwd_fn( + k=k_new, + w1=w_fp16, + w2=h_fp16, + S=S, + BT=BS, + cu_seqlens=cu_seqlens, + ) + q_new_large = transform_q_fwd_fn(q=q_new, w1=w_fp16, w2=h_fp16, cu_seqlens=cu_seqlens, BT=BT, BS=BS, S=S) + w = w.to(q.dtype) + h = h.to(q.dtype) + A = A.to(q.dtype) + dk, dv, _ = parallel_path_bwd_dkv_fn( + q=q_new_large, + k=k_new_large, + v=v, + g_cumsum=g_cumsum, + do=do, + dv=dv, + dg_cumsum=dg_cumsum, + hc_whole=hc_whole, + scale=ctx.scale, + cu_seqlens=cu_seqlens, + L=L, + D=delta, + S=S, + BT=BT, + BS=BS, + ) + dq, dhc_whole, dg_cumsum = parallel_path_bwd_dq_fn( + q=q_new_large, + k=k_new_large, + v=v, + g_cumsum=g_cumsum, + do=do, + dg_cumsum=dg_cumsum, + hc_whole=hc_whole, + scale=ctx.scale, + cu_seqlens=cu_seqlens, + L=L, + D=delta, + S=S, + BT=BT, + BS=BS, + ) + dw1, dw2, dk = chunk_cumprod_householder_bwd_fn( + w1=w, + w2=h, + k=k_new, + dk=dk, + hc_suffix=hc_suffix, + dhc_whole=dhc_whole, + cu_seqlens=cu_seqlens, + S=S, + BT=BS, + ) + dq, dk, dv, dw1, dw2, dg_cumsum = parallel_path_bwd_intra_chunk_fn( + q=q_new, + k=k_new, + v=v, + g_cumsum=g_cumsum, + w1=w, + w2=h, + L=L, + D=delta, + scale=ctx.scale, + dw1=dw1, + dw2=dw2, + dq=dq, + dk=dk, + dv=dv, + do=do, + dg_cumsum=dg_cumsum, + cu_seqlens=cu_seqlens, + S=S, + BT=BS, + ) + dq, dk, dbeta, dw = intra_chunk_preprocess_bwd_fn( + q=q, + k=k, + w=w, + w2=h, + beta=beta, + dq=dq, + dk=dk, + dw1=dw1, + dw2=dw2, + dA_local=dA_local, + A=A, + L=L, + D=delta, + do=do, + scale=ctx.scale, + cu_seqlens=cu_seqlens, + ) + G = q.shape[-2] // k.shape[-2] + if G > 1: + assert dk.dtype == dv.dtype == dw.dtype == dbeta.dtype == torch.float32, 'reduction requires float32' + dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') + dv = reduce(dv, 'b t (h g) k -> b t h k', g=G, reduction='sum') + dw = reduce(dw, 'b t (h g) k -> b t h k', g=G, reduction='sum') + dbeta = reduce(dbeta, 'b t (h g) -> b t h', g=G, reduction='sum') + if dg_cumsum is not None: + dg_cumsum = chunk_global_cumsum(dg_cumsum, cu_seqlens=cu_seqlens, reverse=True) + return (dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dw.to(w.dtype), + dbeta.to(beta.dtype), + dg_cumsum.to(g_cumsum.dtype) if g_cumsum is not None else None, + None, None, None, None) + + +@torch.compiler.disable +def parallel_path_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.Tensor | None = None, + use_cache: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]` + k (torch.Tensor): + keys of shape `[B, T, H, K]` + v (torch.Tensor): + values of shape `[B, T, H, V]` + w (torch.Tensor): + weights of shape `[B, T, H, K]` + beta (torch.Tensor): + beta of shape `[B, T, H]` + g (torch.Tensor): + g of shape `[B, T, HQ]` + scale (float): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + use_cache (bool): + Whether to transform and cache the key values for decoding. Default: `False`. + + Returns: + o (torch.Tensor): + output of shape `[B, T, HQ, V]` + k_cache (torch.Tensor): + k_cache of shape `[B, T, H, K]` + """ + if scale is None: + scale = k.shape[-1]**-0.5 + assert w.dtype == beta.dtype == torch.float32, 'w, beta should be float32 to preserve precision.' + if g is not None: + assert g.dtype == torch.float32, 'g should be float32 to preserve precision.' + assert q.shape[-1] in [16, 32, 64, 128], "only support head_dim in [16, 32, 64, 128] for now. Stay tuned!" + assert v.shape[-1] in [16, 32, 64, 128], "only support head_dim in [16, 32, 64, 128] for now. Stay tuned!" + assert q.shape[-1] == k.shape[-1], 'q, k should have the same head_dim.' + assert k.shape == w.shape, 'k, w should have the same shape.' + assert beta.shape[:3] == k.shape[:3], 'beta should have the same number of heads as k' + if g is not None: + assert g.shape[:3] == q.shape[:3], 'g should have the same number of heads as q' + assert q.shape[-2] % k.shape[-2] == 0, 'the number of query heads should be divisible by the number of key heads' + o, k_cache = ParallelPATHAttentionFunction.apply(q, k, v, w, beta, g, scale, cu_seqlens, use_cache) + return o, k_cache + +parallel_path_attention = parallel_path_attn diff --git a/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py b/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py new file mode 100644 index 0000000000000000000000000000000000000000..12abc9c66f0e0156f3a37e6012efb5c31a433616 --- /dev/null +++ b/fla/ops/path_attn/parallel_path_bwd_inter_dkv.py @@ -0,0 +1,192 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets + + +@triton.heuristics({ + 'USE_GATE': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_path_bwd_dkv_kernel( + q, + k, + v, + g_cumsum, + hc_whole, + scale, + L, + D, + dk, + dv, + do, + dg_cumsum, + cu_seqlens, + indices, + split_offsets, + T, + G: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + S: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATE: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + boh_large = tl.load(split_offsets + i_n).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + boh_large = i_n * tl.cdiv(T, S) + + # offset calculations + + do += (bos * HQ + i_hq) * V + dk += (bos * HQ + i_hq) * K + dv += (bos * HQ + i_hq) * K + L += (bos * HQ + i_hq) + D += (bos * HQ + i_hq) + + k += (bos * H + i_h) * K # GQA when H!=HQ + v += (bos * H + i_h) * V # GQA when H!=HQ + hc_whole += (boh_large * H + i_h) * K * K + + if USE_GATE: + g_cumsum += (bos * HQ + i_hq) + dg_cumsum += (bos * HQ + i_hq) + + # constants + sm_scale = scale * 1.44269504 + + # load query + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + if USE_GATE: + b_g_cumsum_k = tl.zeros([BT], dtype=tl.float32) + p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_g_cumsum_k += tl.load(p_g_cumsum_k, boundary_check=(0, )) + b_dg_cumsum_k = tl.zeros([BT], dtype=tl.float32) + else: + b_g_cumsum_k = None + b_dg_cumsum_k = None + + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + + last_chunk_start = tl.floor(i_t*BT / S).to(tl.int32) * S + idx_j = (tl.floor(i_t * BT / S).to(tl.int32) + 1).to(tl.int32) + + last_chunk_end = tl.ceil(T / BS).to(tl.int32) * BS - BS + + for offset in range(last_chunk_end, last_chunk_start+S-BS, -BS): + p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) + p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) + b_delta = tl.load(p_delta, boundary_check=(0, )) + b_l = tl.load(p_l, boundary_check=(0, )) + + p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j) * HQ + i_hq) * K, (T, K), + (HQ*K*NUM_BLOCKS, 1), (offset, 0), (BS, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_A = tl.dot(b_k, tl.trans(b_q).to(b_k.dtype)) + if USE_GATE: + p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) + b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0, )) + b_A = b_A + b_g_cumsum_q[None, :] - b_g_cumsum_k[:, None] + b_A = tl.where((offset + tl.arange(0, BS) < T)[None, :], b_A, float("-inf")) # avoid nan + b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[None, :]) + p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (offset, 0), (BS, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv += tl.dot(b_A_softmax.to(b_do.dtype), b_do) + b_dp = tl.dot(b_v, tl.trans(b_do)) + + b_dA = ((b_dp - b_delta[None, :]) * b_A_softmax * scale) + if USE_GATE: + b_dg_cumsum_k -= tl.sum(b_dA, axis=1) + b_dk += tl.dot(b_dA.to(b_q.dtype), b_q) + + p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dk, b_dk.to(dk.dtype.element_ty), boundary_check=(0, 1)) + mask = i_t * BT + tl.arange(0, BT) < T + tl.atomic_add( + dv + (i_t * BT + tl.arange(0, BT))[:, None] * HQ * V + tl.arange(0, BV)[None, :], + b_dv, + mask=mask[:, None], + sem='relaxed', + ) + if USE_GATE: + tl.atomic_add(dg_cumsum + (i_t * BT + tl.arange(0, BT)) * HQ, b_dg_cumsum_k, mask=mask, sem='relaxed') + + +def parallel_path_bwd_dkv_fn( + q, k, v, g_cumsum, do, dv, dg_cumsum, + hc_whole, scale, L, D, + cu_seqlens, + S, BT, BS, + chunk_indices: torch.LongTensor | None = None, +): + B, T, num_blocks, HQ, K = q.shape + V = v.shape[-1] + H = k.shape[-2] + G = HQ // H + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + + if cu_seqlens is not None: + assert split_offsets[-1] == hc_whole.shape[0] + + dk = torch.empty(B, T, HQ, K, dtype=torch.float32, device=q.device) + + parallel_path_bwd_dkv_kernel[(NT, B*HQ)]( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + hc_whole=hc_whole, + scale=scale, + L=L, + D=D, + dk=dk, + dv=dv, + do=do, + dg_cumsum=dg_cumsum, + cu_seqlens=cu_seqlens, + indices=indices, + split_offsets=split_offsets, + T=T, + S=S, + BT=BT, + BS=BS, + G=G, + HQ=HQ, + H=H, + K=K, + V=V, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + num_warps=8 if (BT == 128 and K == 128) else 4, + NUM_BLOCKS=num_blocks, + ) + return dk, dv, dg_cumsum diff --git a/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py b/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py new file mode 100644 index 0000000000000000000000000000000000000000..86b51858dc65aa224987b8978ad5f1b29c0186d0 --- /dev/null +++ b/fla/ops/path_attn/parallel_path_bwd_inter_dqh.py @@ -0,0 +1,203 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp2 +from fla.utils import check_shared_mem + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_GATE': lambda args: args['g_cumsum'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_path_bwd_dq_kernel( + q, + k, + v, + g_cumsum, + hc_whole, + scale, + L, + D, + dq, + do, + dhc_whole, + dg_cumsum, + cu_seqlens, + indices, + split_offsets, # varlen specific + T, + G: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + S: tl.constexpr, # aka larger chunk size + NUM_BLOCKS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATE: tl.constexpr, +): + i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + boh_large = tl.load(split_offsets + i_n).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + boh_large = i_n * tl.cdiv(T, S) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + k += (bos * H + i_h) * K # GQA when H!=HQ + v += (bos * H + i_h) * V # GQA when H!=HQ + do += (bos * HQ + i_hq) * V + dq += (bos * HQ + i_hq) * K + hc_whole += (boh_large * H + i_h) * K * K + dhc_whole += (boh_large * HQ + i_hq) * K * K + L += (bos * HQ + i_hq) + D += (bos * HQ + i_hq) + if USE_GATE: + g_cumsum += (bos * HQ + i_hq) + dg_cumsum += (bos * HQ + i_hq) + + # constants + stride_h = H * K * K + stride_hq = HQ * K * K + sm_scale = scale * 1.44269504 + + # load query + p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + p_l = tl.make_block_ptr(L, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_d = tl.make_block_ptr(D, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + b_l = tl.load(p_l, boundary_check=(0,)) + b_delta = tl.load(p_d, boundary_check=(0,)) + + if USE_GATE: + p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)).to(tl.float32) + b_dg_cumsum_q = tl.zeros([BT], dtype=tl.float32) + else: + b_g_cumsum_q = None + b_dg_cumsum_q = None + + curr_end = ((i_t * BT // S) * S).to(tl.int32) + b_dq = tl.zeros([BT, K], dtype=tl.float32) + + for offset_outer in range(0, curr_end, S): + idx_j = offset_outer // S + p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j + 1) * HQ + i_hq) * K, (T, K), + (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + + b_dh = -tl.dot(tl.trans(b_q), b_dq.to(b_q.dtype)) + tl.atomic_add(dhc_whole + idx_j * stride_hq + tl.arange(0, K) + [:, None] * K + tl.arange(0, K)[None, :], b_dh, sem='relaxed') + p_h = tl.make_block_ptr(hc_whole + idx_j * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dq = b_dq - tl.dot(b_dq.to(b_h.dtype), tl.trans(b_h)) + + for offset in range(offset_outer, min(offset_outer+S, i_t*BT), BS): + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (offset, 0), (BS, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A = tl.dot(b_q, tl.trans(b_k).to(b_q.dtype)) + if USE_GATE: + p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (offset,), (BS,), (0,)) + b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)).to(tl.float32) + b_A = b_A + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] + b_A = exp2(b_A * sm_scale - b_l[:, None]) + b_A = tl.where(m_t[:, None], b_A, 0) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (offset, 0), (BS, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dp = tl.dot(b_do, tl.trans(b_v).to(b_do.dtype)) + b_dA = (b_dp - b_delta[:, None]) * b_A * scale + b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) + if USE_GATE: + b_dg_cumsum_q += tl.sum(b_dA, axis=1) + + p_dq = tl.make_block_ptr(dq, (T, K), (K * HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(dq.dtype.element_ty), boundary_check=(0, 1)) + if USE_GATE: + tl.atomic_add(dg_cumsum + o_t * HQ, b_dg_cumsum_q, mask=m_t, sem='relaxed') + + +def parallel_path_bwd_dq_fn( + q, + k, + v, + g_cumsum, + do, + dg_cumsum, + hc_whole, + scale, + L, + D, + cu_seqlens, + S, + BT, + BS, + chunk_indices: torch.LongTensor | None = None, +): + B, T, num_blocks, HQ, K = q.shape + H, V = v.shape[-2:] + G = HQ // H + BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + + # should be NS + if cu_seqlens is not None: + assert split_offsets[-1] == hc_whole.shape[0] + dq = torch.empty(B, T, HQ, K, dtype=torch.float32, device=q.device) + + # [NS, HQ, K, K] instead of [NS, H, K, K] + # atomic add must be initialized to 0 + dhc_whole = torch.zeros(hc_whole.shape[0], HQ, K, K, dtype=torch.float32, device=q.device) + + parallel_path_bwd_dq_kernel[(NT, B*HQ)]( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + hc_whole=hc_whole, + scale=scale, + L=L, + D=D, + dq=dq, + do=do, + dhc_whole=dhc_whole, + dg_cumsum=dg_cumsum, + cu_seqlens=cu_seqlens, + indices=indices, + split_offsets=split_offsets, + T=T, + S=S, + BT=BT, + BS=BS, + G=G, + HQ=HQ, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + NUM_BLOCKS=num_blocks, + num_warps=8 if (BT == 128 and K == 128) else 4, + num_stages=3 if check_shared_mem('ampere') else 2, + ) + return dq, dhc_whole, dg_cumsum diff --git a/fla/ops/path_attn/parallel_path_bwd_intra.py b/fla/ops/path_attn/parallel_path_bwd_intra.py new file mode 100644 index 0000000000000000000000000000000000000000..220c2d03be757fcb139afe142c353e751d857f1f --- /dev/null +++ b/fla/ops/path_attn/parallel_path_bwd_intra.py @@ -0,0 +1,175 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['offsets'] is not None, + 'USE_GATE': lambda args: args['g_cumsum'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_path_bwd_intra_chunk_kernel( + q, k, v, g_cumsum, w1, w2, + L, D, + dq, dq_new, dk, dv, dw1, dw2, do, dg_cumsum, + offsets, indices, + T, scale, + G: tl.constexpr, HQ: tl.constexpr, H: tl.constexpr, + K: tl.constexpr, V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + BT: tl.constexpr, S: tl.constexpr, + IS_VARLEN: tl.constexpr, USE_GATE: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + # offset calculations + k += (bos * H + i_h) * K # GQA when H!=HQ + v += (bos * H + i_h) * V # GQA when H!=HQ + w1 += (bos * H + i_h) * K + w2 += (bos * H + i_h) * K + + q += (bos * HQ + i_hq) * K + dq += (bos * HQ + i_hq) * K + dq_new += (bos * HQ + i_hq) * K + dk += (bos * HQ + i_hq) * K + dv += (bos * HQ + i_hq) * V + do += (bos * HQ + i_hq) * V + dw1 += (bos * HQ + i_hq) * K + dw2 += (bos * HQ + i_hq) * K + L += (bos * HQ + i_hq) + D += (bos * HQ + i_hq) + if USE_GATE: + g_cumsum += (bos * HQ + i_hq) + dg_cumsum += (bos * HQ + i_hq) + + # constants + sm_scale = scale * 1.44269504 + + p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_delta = tl.load(p_delta, boundary_check=(0, )) + p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_l = tl.load(p_l, boundary_check=(0, )) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + p_dq = tl.make_block_ptr(dq, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_dq += tl.load(p_dq, boundary_check=(0, 1)) + p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + + if USE_GATE: + p_gq_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) + b_gq_cumsum = tl.load(p_gq_cumsum, boundary_check=(0, )) + b_dgq = tl.zeros([BT], dtype=tl.float32) + else: + b_dgq = None + + curr_start = (tl.floor(i_t * BT / S).to(tl.int32) * S).to(tl.int32) + + for offset in range(curr_start, i_t * BT, BT): + mask = offset + tl.arange(0, BT) < T + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q_tmp = tl.zeros([BT, BK], dtype=tl.float32) + b_q_tmp += b_q + for i_t_small in range(i_t * BT - BT, offset, -BT): + p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_A_tmp = tl.dot(b_q_tmp.to(b_w1.dtype), tl.trans(b_w1)) + b_q_tmp -= tl.dot(b_A_tmp.to(b_w1.dtype), b_w2) + b_q2 = b_q_tmp.to(b_k.dtype) + b_A = tl.dot(b_q2, tl.trans(b_k)) + if USE_GATE: + p_gk_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BT, ), (0, )) + b_gk_cumsum = tl.load(p_gk_cumsum, boundary_check=(0, )) + b_A = b_A + b_gq_cumsum[:, None] - b_gk_cumsum[None, :] + b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan + b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[:, None]) + b_dv = tl.dot(tl.trans(b_A_softmax.to(b_do.dtype)), b_do) + tl.atomic_add( + dv + ((offset + tl.arange(0, BT)) * HQ * V)[:, None] + tl.arange(0, BV)[None, :], + b_dv.to(dv.dtype.element_ty), + mask=mask[:, None], + sem='relaxed', + ) + p_v = tl.make_block_ptr(v, (T, V), (V*H, 1), (offset, 0), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dp = tl.dot(b_do, tl.trans(b_v)) + b_dA = ((b_dp - b_delta[:, None]) * b_A_softmax * scale) + if USE_GATE: + b_dgk = -tl.sum(b_dA, axis=0) + tl.atomic_add(dg_cumsum + (offset + tl.arange(0, BT)) * HQ, b_dgk, mask=mask, sem='relaxed') + b_dgq += tl.sum(b_dA, axis=1) + b_dA = b_dA.to(b_v.dtype) + b_dk = tl.dot(tl.trans(b_dA), b_q2) + tl.atomic_add(dk + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, + BK)[None, :], b_dk, mask=mask[:, None], sem='relaxed') + p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_dA2 = tl.dot(b_dq.to(b_w2.dtype), tl.trans(b_w2)).to(b_v.dtype) + b_A2 = tl.dot(b_q2.to(b_w1.dtype), tl.trans(b_w1)).to(b_v.dtype) + b_dw2 = -tl.dot(tl.trans(b_A2), b_dq.to(b_v.dtype)) + tl.atomic_add(dw2 + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, + BK)[None, :], b_dw2, mask=mask[:, None], sem='relaxed') + b_dw1 = -tl.dot(tl.trans(b_dA2), b_q2.to(b_v.dtype)) + tl.atomic_add(dw1 + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, + BK)[None, :], b_dw1, mask=mask[:, None], sem='relaxed') + b_dq -= tl.dot(b_dA2, b_w1.to(b_v.dtype)) + b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) + + p_dq_new = tl.make_block_ptr(dq_new, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_dq_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) + mask = i_t * BT + tl.arange(0, BT) < T + if USE_GATE: + tl.atomic_add(dg_cumsum + (i_t * BT + tl.arange(0, BT)) * HQ, b_dgq, mask=mask, sem='relaxed') + + +def parallel_path_bwd_intra_chunk_fn( + q, k, v, g_cumsum, w1, w2, + dq, dk, dv, dg_cumsum, dw1, dw2, do, + scale, L, D, + cu_seqlens, + S, BT, + chunk_indices: torch.LongTensor | None = None, +): + assert dk.dtype == dv.dtype == dw1.dtype == dw2.dtype == torch.float32, 'atomic_add requires float32' + B, T, HQ, K = q.shape + assert dk.shape == dq.shape + + V = v.shape[-1] + H = k.shape[-2] + G = HQ // H + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + dq_new = torch.empty_like(dq, dtype=q.dtype) + parallel_path_bwd_intra_chunk_kernel[(NT, B*HQ)]( + q=q, k=k, v=v, g_cumsum=g_cumsum, + w1=w1, w2=w2, L=L, D=D, + dq=dq, dq_new=dq_new, dk=dk, dv=dv, dw1=dw1, dw2=dw2, + do=do, dg_cumsum=dg_cumsum, + offsets=cu_seqlens, indices=indices, + T=T, S=S, BT=BT, scale=scale, + G=G, HQ=HQ, H=H, K=K, V=V, + BK=triton.next_power_of_2(K), BV=triton.next_power_of_2(V), + ) + return dq_new, dk, dv, dw1, dw2, dg_cumsum diff --git a/fla/ops/path_attn/parallel_path_fwd.py b/fla/ops/path_attn/parallel_path_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..a5133bf76f1e3d93b0199b9e8ded36195b581baf --- /dev/null +++ b/fla/ops/path_attn/parallel_path_fwd.py @@ -0,0 +1,198 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices + + +@triton.heuristics({ + 'USE_GATE': lambda args: args['g_cumsum'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_path_fwd_kernel( + q, + k, + v, + o, + o_new, + g_cumsum, + w1, + w2, + scale, + L, + L_new, + M, + cu_seqlens, + indices, + T, + G: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_GATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_q = tl.zeros([BT, BK], dtype=tl.float32) + b_q += tl.load(p_q, boundary_check=(0, 1)) + sm_scale = scale * 1.44269504 + b_o = tl.zeros([BT, BV], dtype=tl.float32) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + + p_L = tl.make_block_ptr(L + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) + p_M = tl.make_block_ptr(M + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) + b_l = tl.load(p_L, boundary_check=(0,)) + b_m = tl.load(p_M, boundary_check=(0,)) + + if USE_GATE: + p_g_cumsum_q = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) + b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)) + else: + b_g_cumsum_q = None + + for offset in range((i_t + 1) * BT - 2 * BS, i_t*BT-BS, -BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ + p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) + p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) + # [BK, BS] + + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BK, BK] + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + # [BT, BS] + m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) + b_s = tl.dot(b_q.to(b_k.dtype), b_k) + + if USE_GATE: + p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) + b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) + b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] + b_s = tl.where(m_s[:, None], b_s * sm_scale, float("-inf")) + b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) + alpha = tl.math.exp2(b_m - b_m_new) + b_s = tl.math.exp2(b_s - b_m_new[:, None]) + b_o *= alpha[:, None] + b_l = b_l * alpha + tl.sum(b_s, 1) + b_m = b_m_new + b_o += tl.dot(b_s.to(b_v.dtype), b_v) + b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) + b_s2 = tl.where(m_s[:, None], b_s2, 0) + b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) + + tl.debug_barrier() + + for offset in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ + p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) + p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + # [BT, BS] + b_s = tl.dot(b_q.to(b_k.dtype), b_k) + if USE_GATE: + p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) + b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) + b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] + b_s = b_s * sm_scale + b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) + alpha = tl.math.exp2(b_m - b_m_new) + b_s = tl.math.exp2(b_s - b_m_new[:, None]) + b_o *= alpha[:, None] + b_l = b_l * alpha + tl.sum(b_s, 1) + b_m = b_m_new + b_o += tl.dot(b_s.to(b_v.dtype), b_v) + b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) + b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) + + b_o = b_o / b_l[:, None] + p_o_new = tl.make_block_ptr(o_new + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + tl.store(p_o_new, b_o.to(p_o_new.dtype.element_ty), boundary_check=(0, 1)) + b_l = tl.math.log2(b_l) + b_m + p_L_new = tl.make_block_ptr(L_new + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) + tl.store(p_L_new, b_l.to(p_L_new.dtype.element_ty), boundary_check=(0,)) + + +def parallel_path_fwd_fn( + q, + k, + v, + o, + g_cumsum, + w1, + w2, + scale, + L, + M, + cu_seqlens, + BT, + BS, + chunk_indices: torch.LongTensor | None = None, +): + B, T, HQ, K = q.shape + V = v.shape[-1] + H = k.shape[-2] + G = HQ // H + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + grid = (NT, B * HQ) + o_new = torch.empty_like(o, dtype=v.dtype) + L_new = torch.empty_like(L) + + parallel_path_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + o_new=o_new, + w1=w1, + w2=w2, + g_cumsum=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + indices=indices, + L=L, + L_new=L_new, + M=M, + T=T, + K=K, + V=V, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + G=G, + HQ=HQ, + H=H, + BS=BS, + BT=BT, + num_warps=8 if (BT == 128 and K == 128) else 4, + ) + return o_new, L_new diff --git a/fla/ops/path_attn/prepare_k_cache.py b/fla/ops/path_attn/prepare_k_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..bc8708953180af5d130fd824826ec77169a27bb7 --- /dev/null +++ b/fla/ops/path_attn/prepare_k_cache.py @@ -0,0 +1,76 @@ +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['offsets'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def parallel_path_fwd_kernel_prepare_k_cache( + k, k_new, w1, w2, + offsets, indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + k += (bos * H + i_h) * K + k_new += (bos * H + i_h) * K + w1 += (bos * H + i_h) * K + w2 += (bos * H + i_h) * K + # constants + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_k = tl.zeros([BT, BK], dtype=tl.float32) + b_k += tl.load(p_k, boundary_check=(0, 1)) + for k_block_idx in range(i_t + 1, tl.cdiv(T, BT)): + p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) + p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_A = tl.dot(b_k.to(b_w2.dtype), tl.trans(b_w2)) + b_k = b_k - tl.dot(b_A.to(b_w1.dtype), b_w1) + + p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + + +def prepare_k_cache_fn(k, w1, w2, cu_seqlens, BS, use_cache=False, chunk_indices: torch.LongTensor | None = None): + if not use_cache: + return None + else: + B, T, H, K = k.shape + k_new = torch.empty_like(k) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BS) + indices = chunk_indices + NT = triton.cdiv(T, BS) if cu_seqlens is None else len(indices) + grid = (NT, B * H) + parallel_path_fwd_kernel_prepare_k_cache[grid]( + k=k, + k_new=k_new, + w1=w1, + w2=w2, + offsets=cu_seqlens, + indices=indices, + H=H, + T=T, + K=K, + BT=BS, + BK=triton.next_power_of_2(K), + ) + return k_new diff --git a/fla/ops/path_attn/transform_q.py b/fla/ops/path_attn/transform_q.py new file mode 100644 index 0000000000000000000000000000000000000000..3edbffdaa6308b4d17428e4796add94d4c827b9b --- /dev/null +++ b/fla/ops/path_attn/transform_q.py @@ -0,0 +1,110 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import get_max_num_splits, prepare_chunk_indices + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def transform_q_fwd_kernel( + q, + q_new, + w1, + w2, + cu_seqlens, + indices, + T, + S: tl.constexpr, + G: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if IS_VARLEN: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + # boh = i_n * tl.cdiv(T, BS) + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_q = tl.zeros([BT, BK], dtype=tl.float32) + b_q += tl.load(p_q, boundary_check=(0, 1)) + + if BS == BT: + if (i_t * BT) % S == 0: + p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (i_t * BT // S)) * HQ + i_hq) * K, + (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) + + for offset in range((i_t + 1) * BT - 2 * BS, S-BS, -BS): + p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) + p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) + b_w1 = tl.load(p_w1, boundary_check=(0, 1)) + b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) + b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) + b_s2 = tl.where(m_s[:, None], b_s2, 0) + b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) + + if offset % S == 0: + p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (offset // S)) * HQ + i_hq) * K, + (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) + + +def transform_q_fwd_fn( + q, + w1, + w2, + cu_seqlens, + BT, + BS, + S, + chunk_indices: torch.LongTensor | None = None, +): + B, T, HQ, K = q.shape + H = w1.shape[-2] + G = HQ // H + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + indices = chunk_indices + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) + + num_blocks = triton.cdiv(T, S) if cu_seqlens is None else get_max_num_splits(cu_seqlens, S) + q_new = torch.zeros(B, T, num_blocks, HQ, K, dtype=q.dtype, device=q.device) + transform_q_fwd_kernel[(NT, B * HQ)]( + q=q, + q_new=q_new, + w1=w1, + w2=w2, + cu_seqlens=cu_seqlens, + indices=indices, + T=T, + K=K, + BK=triton.next_power_of_2(K), + G=G, + HQ=HQ, + H=H, + BS=BS, + BT=BT, + S=S, + NUM_BLOCKS=num_blocks, + num_warps=8 if (BT == 128 and K == 128) else 4, + ) + return q_new diff --git a/fla/ops/quasar/__init__.py b/fla/ops/quasar/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8788c611d75aefd62ec9e0d056457e3ddaf14cef --- /dev/null +++ b/fla/ops/quasar/__init__.py @@ -0,0 +1,4 @@ +from .chunk import chunk_quasar +from .fused_recurrent import fused_recurrent_quasar + +__all__ = ['chunk_quasar', 'fused_recurrent_quasar'] diff --git a/fla/ops/quasar/chunk.py b/fla/ops/quasar/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..533ee79d2d5f2ddb92513bc2ea10d8f057d319b3 --- /dev/null +++ b/fla/ops/quasar/chunk.py @@ -0,0 +1,372 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention + +import torch +import triton + +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from fla.ops.gla.chunk import chunk_gla_fwd_o_gk +from fla.ops.quasar.chunk_intra import chunk_quasar_fwd_intra +from fla.ops.quasar.gate import fused_quasar_gate, fast_quasar_alpha +from fla.utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard +from fla.ops.common.chunk_o import chunk_fwd_o, chunk_bwd_dv_local, chunk_bwd_dqkwg + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +@input_guard +def chunk_quasar_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = 64, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Kernelized chunk-wise QuasarAttention forward pass.""" + B, T, H, S = q.shape + BT = chunk_size + if BT != 64: + raise ValueError("Only chunk_size=64 is currently supported in the kernelized Quasar chunk path") + + # Prepare chunk indices for varlen + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + + # Quasar-specific per-token alpha + # alpha[t] = (1 - exp(-beta * ||k_t||^2)) / (||k_t||^2 + eps) + # beta is head-wise [H] + + # Ensure high precision for stability components + k_f32 = k.float() + k_norm_sq = (k_f32 * k_f32).sum(dim=-1) # [B, T, H] + + # Aggressive clamping to prevent exp() instability + k_norm_sq = torch.clamp(k_norm_sq, min=0.1, max=10.0) + + # Flexible beta shape: support head-wise [H] or token-wise [B, T, H] + if beta.dim() == 1: + beta_h = beta.view(1, 1, H).float() + else: + beta_h = beta.float() + + # Quasar-style decay computation with per-dim dt_bias + # dt_bias is [H*K], we keep it full dimensional like Quasar does + + if A_log is not None: + A = A_log.float().exp().view(1, 1, H, 1) # [1, 1, H, 1] for broadcasting + else: + A = 1.0 + + # Expand beta to [B, T, H, 1] to match key dim + beta_expanded = beta_h.unsqueeze(-1) # [B, T, H, 1] + + # Reshape dt_bias to [H, K] and add batch/time dims + if dt_bias is not None: + K = q.shape[-1] # key dimension + dt_bias_full = dt_bias.float().view(1, 1, H, K) # [1, 1, H, K] + else: + dt_bias_full = 0.0 + K = q.shape[-1] + + # Expand k_norm_sq to [B, T, H, 1] for broadcasting + k_norm_sq_expanded = k_norm_sq.unsqueeze(-1) # [B, T, H, 1] + + # Compute Quasar-style gate per-dimension: -exp(A_log) * softplus(beta + dt_bias) + g_quasar = -A * torch.nn.functional.softplus(beta_expanded + dt_bias_full) # [B, T, H, K] + + # Convert to decay factor + decay = torch.exp(g_quasar) # [B, T, H, K] + + # Quasar alpha formula adapted per-dimension + alpha = (1.0 - decay) / (k_norm_sq_expanded + 1e-6) # [B, T, H, K] + + # For Quasar's kernel which expects beta_tok as [B, T, H], we take mean across K + # This is a compromise - ideally the kernel would handle per-dim + beta_tok = alpha.mean(dim=-1).clamp_(min=1e-4, max=0.95).to(dtype=q.dtype) # [B, T, H] + + # Use a zero decay tensor to reuse kernels without additional gating. + # Shape-compatible with log-space decay, but equals 0 -> exp(0)=1. + g_zero = torch.zeros_like(q) + + scale = S ** -0.5 + + # Intra-chunk: compute Aqk + Akk^{-1} representation and WY factors (w/u). + w, u, qg, kg, Aqk, Akk = chunk_quasar_fwd_intra( + q=q, + k=k, + v=v, + gk=g_zero, + beta=beta_tok, # FIXED: pass per-token alpha, not head-wise beta + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + safe_gate=True, + disable_recompute=True, + beta_out=beta_tok, # Output is same as input for Quasar + ) + + # Recurrence (kernelized, no Python loop): produces per-chunk states h and updated values v_new. + if initial_state is not None and initial_state.dtype != torch.float32: + initial_state_f32 = initial_state.float() + else: + initial_state_f32 = initial_state + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + g=None, + gk=None, + initial_state=initial_state_f32, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + + # Output (kernelized): o = q @ h + Aqk @ v_new (implemented via efficient SRAM standard kernel) + o = chunk_fwd_o( + q=q, + k=kg, # standard k was normalized, use scaled kg here + v=v_new, + h=h, + g=None, + scale=scale, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + +class ChunkQuasarFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, + ): + chunk_size = 64 + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size) if cu_seqlens is not None else None + + o, final_state = chunk_quasar_fwd( + q=q, + k=k, + v=v, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + ctx.save_for_backward(q, k, v, beta, A_log, dt_bias, initial_state, cu_seqlens, chunk_indices) + ctx.chunk_size = chunk_size + ctx.output_final_state = output_final_state + + return o, final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do: torch.Tensor, d_final_state: torch.Tensor | None): + q, k, v, beta, A_log, dt_bias, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors + chunk_size = ctx.chunk_size + + # Recompute forward intermediates (simpler than saving all) + B, T, H, S = q.shape + + # Recompute alpha + eps = 1e-6 + k_norm_sq = (k.float() * k.float()).sum(dim=-1) # [B, T, H] + k_norm_sq = torch.clamp(k_norm_sq, min=0.1, max=10.0) + + if beta.dim() == 1: + beta_h = beta.view(1, 1, H).to(k_norm_sq.dtype) + else: + beta_h = beta.to(k_norm_sq.dtype) + + beta_h = torch.clamp(beta_h, min=0.01, max=10.0) + # Compute alpha with numerical stability + exp_term = torch.exp(-beta_h * k_norm_sq) + alpha = (1.0 - exp_term) / (k_norm_sq + eps) + beta_tok = alpha.clamp_(min=1e-4, max=0.95).to(dtype=q.dtype) + + g_zero = torch.zeros_like(q) + scale = S ** -0.5 + + # Allocate beta_out for Quasar alpha computation + beta_out = torch.empty_like(beta_tok) + + # Recompute forward intermediates + w, u, qg, kg, Aqk, Akk = chunk_quasar_fwd_intra( + q=q, + k=k, + v=v, + gk=g_zero, + beta=beta_tok, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + safe_gate=False, + disable_recompute=True, + beta_out=beta_out, + ) + + if initial_state is not None and initial_state.dtype != torch.float32: + initial_state_f32 = initial_state.float() + else: + initial_state_f32 = initial_state + + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + g=None, + gk=None, + initial_state=initial_state_f32, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + # Backward: output kernel (dA, dv) + from fla.ops.quasar.chunk_bwd import chunk_quasar_bwd_dAv + dA, dv = chunk_quasar_bwd_dAv( + q=q, + k=k, + v=v_new, + do=do, + A=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # Backward: recurrence (dh) + from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu + dh, dh0, dv2 = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=kg, + w=w, + do=do, + dv=dv, + g=None, + gk=None, + h0=initial_state_f32, + dht=None, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + use_exp2=True, + ) + dv = dv2 + + # Backward: WY recompute + intra (dq, dk, dbeta) + from fla.ops.quasar.chunk_bwd import chunk_quasar_bwd_wy_dqkb_fused + dq, dk, dv3, db, dA2 = chunk_quasar_bwd_wy_dqkb_fused( + q=q, + k=k, + v=v, + v_new=v_new, + beta=beta_tok, + A=Akk, + h=h, + do=do, + dh=dh, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # Combine gradients + dv = dv + dv3 + dA = dA + dA2 + + # Backward: alpha formula (dbeta from dk) + # db is the gradient of Loss w.r.t alpha, shape [B, T, H] + db_f32 = db.float() + + # Aggressive clamping for gradient stability + k_norm_sq = torch.clamp(k_norm_sq, min=0.1, max=10.0) + + if beta.dim() == 1: + beta_h = beta.view(1, 1, H).float() + beta_h = torch.clamp(beta_h, min=0.01, max=10.0) + # Chain rule: dL/dbeta_head = sum( dL/dalpha * dalpha/dbeta ) + dalpha_dbeta = k_norm_sq * exp_term / (k_norm_sq + eps) + dbeta = (db_f32 * dalpha_dbeta).sum(dim=(0, 1)) / T + dbeta = torch.clamp(dbeta, min=-1.0, max=1.0) + else: + beta_h = beta.float() + beta_h = torch.clamp(beta_h, min=0.01, max=10.0) + # Chain rule: dL/dbeta_token = dL/dalpha * dalpha/dbeta + dalpha_dbeta = k_norm_sq * exp_term / (k_norm_sq + eps) + dbeta = db_f32 * dalpha_dbeta + # Token-wise gradient doesn't need / T normalization if it's fed to linear layer + dbeta = torch.clamp(dbeta, min=-1.0, max=1.0) + + return dq, dk, dv, dbeta, None, None, None, None, None + + +@torch.compiler.disable +def chunk_quasar( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Chunk-wise QuasarAttention forward pass with autograd support. + + Args: + q (torch.Tensor): Query tensor of shape [B, T, H, S] + k (torch.Tensor): Key tensor of shape [B, T, H, S] + v (torch.Tensor): Value tensor of shape [B, T, H, S] + beta (torch.Tensor): Beta parameter tensor of shape [H] + A_log (torch.Tensor | None): Learnable state decay, shape [H] + dt_bias (torch.Tensor | None): Learnable time bias, shape [H*K] + initial_state (torch.Tensor | None): Initial state tensor of shape [B, H, S, S] + output_final_state (bool): Whether to output the final state + cu_seqlens (torch.Tensor | None): Cumulative sequence lengths for variable-length sequences + + Returns: + o (torch.Tensor): Output tensor of shape [B, T, H, S] + final_state (torch.Tensor | None): Final state tensor of shape [B, H, S, S] if output_final_state + """ + return ChunkQuasarFunction.apply(q, k, v, beta, A_log, dt_bias, initial_state, output_final_state, cu_seqlens) diff --git a/fla/ops/quasar/chunk_bwd.py b/fla/ops/quasar/chunk_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..ec6b1f367d527fc4b9b6dcfff2aa88b712172759 --- /dev/null +++ b/fla/ops/quasar/chunk_bwd.py @@ -0,0 +1,362 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention (no gating) + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_NVIDIA_HOPPER, IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem + +@triton.jit +def safe_dot(a, b): + return tl.inline_asm_elementwise( + asm="mov.f32 $0, $1;", + constraints="=r,r", + args=[tl.dot(a, b)], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_quasar_bwd_kernel_dAv( + q, + k, + v, + A, + do, + dv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + q += (bos * H + i_h).to(tl.int64) * K + k += (bos * H + i_h).to(tl.int64) * K + v += (bos * H + i_h).to(tl.int64) * V + do += (bos * H + i_h).to(tl.int64) * V + dv += (bos * H + i_h).to(tl.int64) * V + dA += (bos * H + i_h).to(tl.int64) * BT + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dA += safe_dot(b_do, b_v) + b_dv = safe_dot(b_A.to(b_do.dtype), b_do) + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_quasar_bwd_kernel_wy_dqkb_fused( + q, + k, + v, + v_new, + beta, + A, + h, + do, + dh, + dq, + dk, + dv, + dv2, + db, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b = i_bh // H + i_h = i_bh % H + + if IS_VARLEN: + i_tg = i_t.to(tl.int64) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + beta += (bos * H + i_h) + A += (bos * H + i_h) * BT + h += (i_tg * H + i_h) * K*V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K*V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dv2 += (bos * H + i_h) * V + db += (bos * H + i_h) + dA += (bos * H + i_h) * BT + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_db = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v_new = tl.make_block_ptr(v_new, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BT, BV] + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) + b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) + tl.debug_barrier() # DO NOT REMOVE THIS LINE! + if i_k == 0: + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv_inner = tl.load(p_dv, boundary_check=(0, 1)) # Re-load to break dependency across debug_barrier + + b_dA += tl.dot(b_dv_inner.to(b_v.dtype), tl.trans(b_v)) + + b_dvb = tl.dot(b_A, b_dv_inner.to(b_A.dtype)) + b_dv2 = b_dvb * b_beta[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = b_dq * scale + b_dw = -b_dw.to(b_A.dtype) + b_dA += tl.dot(b_dw, tl.trans(b_k.to(b_A.dtype))) + + b_dkgb = tl.dot(b_A, b_dw) + + b_db += tl.sum(b_dkgb * b_k, 1) + + b_dk = b_dk + b_dkgb * b_beta[:, None] + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + b_dA = tl.where(m_A, -b_dA, 0) + + p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def chunk_quasar_bwd_dAv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + A: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem: + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dA = v.new_empty(B, T, H, BT, dtype=torch.float) + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_quasar_bwd_kernel_dAv[grid]( + q=q, + k=k, + v=v, + A=A, + do=do, + dv=dv, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dA, dv + + +def chunk_quasar_bwd_wy_dqkb_fused( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dv2 = torch.empty_like(v) + db = q.new_empty(B, T, H, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + # H200/B200 Triton autotune can benchmark invalid BK/BV/warps candidates for + # this Quasar backward and leave the CUDA context in illegal-address state. + # Force a conservative tile that fits head_dim=128 without autotune. + BK = min(max(triton.next_power_of_2(K), 16), 32) + BV = min(max(triton.next_power_of_2(V), 16), 32) + + grid = (NT, B * H) + chunk_quasar_bwd_kernel_wy_dqkb_fused[grid]( + q=q, + k=k, + v=v, + v_new=v_new, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dv=dv, + dv2=dv2, + db=db, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + num_warps=4, + num_stages=3, + ) + dv = dv2 + return dq, dk, dv, db, dA diff --git a/fla/ops/quasar/chunk_intra.py b/fla/ops/quasar/chunk_intra.py new file mode 100644 index 0000000000000000000000000000000000000000..6ddf2277a4993e64e4aefcec238b689584ba4d9d --- /dev/null +++ b/fla/ops/quasar/chunk_intra.py @@ -0,0 +1,903 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.quasar.wy_fast import recompute_w_u_fwd +from fla.ops.quasar.chunk_intra_token_parallel import chunk_quasar_fwd_intra_token_parallel +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2, gather +from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs + +if IS_TF32_SUPPORTED: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') +else: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "K", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_quasar_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + Aqk += (bos * H + i_h) * BT + Akk += (bos * H + i_h) * BT + Akkd += (bos * H + i_h) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * H*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b1 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b2 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b3 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * H*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_quasar_bwd_kernel_intra( + q, + k, + g, + beta, + dAqk, + dAkk, + dq, + dq2, + dk, + dk2, + dg, + dg2, + db, + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, + SAFE_GATE: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_k, i_i = i_kc // NC, i_kc % NC + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + beta += bos * H + i_h + + dAqk += (bos * H + i_h) * BT + dAkk += (bos * H + i_h) * BT + dq += (bos * H + i_h) * K + dq2 += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dk2 += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + dg2 += (bos * H + i_h) * K + db += (i_k * all + bos) * H + i_h + + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) + b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) + if i_i > 0: + p_gn = g + i_ti * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp2(b_gn - b_gk) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + # [BC, BK] + b_dq2 += tl.dot(b_dAqk, b_kg) + b_dk2 += tl.dot(b_dAkk, b_kg) + b_gqn = exp2(b_g - b_gn) + b_dq2 *= b_gqn + b_dk2 *= b_gqn + + o_i = tl.arange(0, BC) + m_dA = (i_ti + o_i) < T + o_dA = (i_ti + o_i) * H*BT + i_i * BC + p_kj = k + i_ti * H*K + o_k + p_gkj = g + i_ti * H*K + o_k + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_qk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) + b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) + b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) + exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) + exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) + + b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) + b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) + # [BK] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_gqk = exp2(b_g - b_gkj[None, :]) + b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * b_gqk, 0.) + b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.) + + p_kj += H*K + p_gkj += H*K + + b_db = tl.sum(b_dk2 * b_k, 1) + b_dk2 *= b_b[:, None] + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_ti,), (BC,), (0,)) + + b_dg2 = b_q * b_dq2 + b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + tl.debug_barrier() + b_dkt = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = g + (min(i_ti + BC, T) - 1) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + for i_j in range(i_i + 1, NC): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + # [BC] + b_b = tl.load(p_b, boundary_check=(0,)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC, BK] + b_gkn = exp2(b_gk - b_gn) + b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) + b_kbg = b_kb * tl.where(m_j[:, None], b_gkn, 0) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dkt += tl.dot(b_dAqk, b_qg) + b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt *= exp2(b_gn - b_g) + o_dA = i_ti * H*BT + i_i * BC + o_i + p_qj = q + i_ti * H*K + o_k + p_kj = k + i_ti * H*K + o_k + p_gkj = g + i_ti * H*K + o_k + p_bj = beta + i_ti * H + + if SAFE_GATE: + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) + b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + + m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_kk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) + b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) + # ensure numerical stability + b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) + exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) + exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) + + b_q_exp = b_q * exp_b_g_diag_kk + b_kb_exp = b_k * b_b[:, None] * exp_b_g_diag_kk + + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dAqk = tl.load(dAqk + o_dA + j * H*BT) + b_dAkk = tl.load(dAkk + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_gkq = exp2(b_gkj[None, :] - b_g) + b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.) + b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * b_gkq, 0.) + + p_qj += H*K + p_kj += H*K + p_gkj += H*K + p_bj += H + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dg2 = tl.make_block_ptr(dg2, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + + b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) + b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) + b_dk2 += b_dkt + + tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_quasar_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * H + i_h) * K + beta = beta + bos * H + i_h + Aqk = Aqk + (bos * H + i_h) * BT + Akk = Akk + (bos * H + i_h) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (H*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * H*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_quasar_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, + chunk_indices: tuple[torch.Tensor] | None = None, + safe_gate: bool = False, + disable_recompute: bool = False, + beta_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K = k.shape + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) + + # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) + # Step 1: compute diagonal blocks into Akk_diag (fp32) + if safe_gate: + grid = (NT, NC, B * H) + BK = triton.next_power_of_2(K) + chunk_quasar_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + else: + Aqk, Akkd = chunk_quasar_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + sub_chunk_size=sub_chunk_size, + beta_out=beta_out, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + grid = (NT, B * H) + chunk_quasar_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + ) + w, u, qg, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=Akk, + q=q if disable_recompute else None, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, qg, kg, Aqk, Akk + + +def chunk_quasar_bwd_intra( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, +): + B, T, H, K = k.shape + BT = chunk_size + BC = min(16, BT) + BK = min(32, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq2 = torch.empty_like(q) + dk2 = torch.empty_like(k) + db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) + dg2 = torch.empty_like(dg, dtype=torch.float) + grid = (NK * NC, NT, B * H) + chunk_quasar_bwd_kernel_intra[grid]( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dq2=dq2, + dk=dk, + dk2=dk2, + dg=dg, + dg2=dg2, + db=db2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + SAFE_GATE=safe_gate, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + dq = dq2 + dk = dk2 + db = db2.sum(0).add_(db) + dg = dg2 + + return dq, dk, db, dg diff --git a/fla/ops/quasar/chunk_intra_token_parallel.py b/fla/ops/quasar/chunk_intra_token_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..0b84deadc3d766bb2b452c18c3c75c5e5218abd3 --- /dev/null +++ b/fla/ops/quasar/chunk_intra_token_parallel.py @@ -0,0 +1,187 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp2 +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_QUASAR_ALPHA': lambda args: args['beta_out'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_quasar_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + beta_out, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_QUASAR_ALPHA: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + q += bos * H*K + k += bos * H*K + g += bos * H*K + Aqk += bos * H*BT + Akk += bos * H*BC + beta += bos * H + if USE_QUASAR_ALPHA: + beta_out += bos * H + + BK: tl.constexpr = triton.next_power_of_2(K) + o_h = tl.arange(0, BH) + o_k = tl.arange(0, BK) + m_h = (i_hg * BH + o_h) < H + m_k = o_k < K + + p_q = tl.make_block_ptr(q + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_g = tl.make_block_ptr(g + i_t * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_t * H, (H,), (1,), (i_hg * BH,), (BH,), (0,)) + # [BH, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + # 1. QUASAR CT ALGORITHM + b_k2 = tl.sum(b_k * b_k, axis=1) + b_k2_clamped = tl.where(b_k2 < 0.1, 0.1, tl.where(b_k2 > 10.0, 10.0, b_k2)) + b_alpha = (1.0 - tl.exp(-b_beta * b_k2_clamped)) / (b_k2_clamped + 1e-6) + + p_beta_out = tl.make_block_ptr(beta_out + i_t * H, (H,), (1,), (i_hg * BH,), (BH,), (0,)) + tl.store(p_beta_out, b_alpha.to(beta_out.dtype.element_ty), boundary_check=(0,)) + + b_k = b_k * b_alpha[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + p_kj = tl.make_block_ptr(k + j * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_gj = tl.make_block_ptr(g + j * H*K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + # [BH, BK] + b_kj = tl.load(p_kj, boundary_check=(0, 1)).to(tl.float32) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = b_kj * exp2(b_g - b_gj) + + b_kgj = tl.where(m_k[None, :], b_kgj, 0.0) + # [BH] + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * H*BT + (i_hg * BH + o_h) * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_h) + tl.store(Akk + i_t * H*BC + (i_hg * BH + o_h) * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_h) + + +def chunk_quasar_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, + beta_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, H, K] cumsum of gates + beta: [B, T, H] + Aqk: [B, T, H, BT] output tensor to write to + Akk: [B, T, H, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + beta_out: Required output tensor for Quasar CT scalar decay + """ + B, T, H, K = q.shape + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(H, meta['BH'])) + chunk_quasar_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + beta_out=beta_out, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/fla/ops/quasar/forward_substitution.py b/fla/ops/quasar/forward_substitution.py new file mode 100644 index 0000000000000000000000000000000000000000..33adeb725a01fc0387d40f79952ed9a702fd6cec --- /dev/null +++ b/fla/ops/quasar/forward_substitution.py @@ -0,0 +1,135 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention + +import torch +import triton +import triton.language as tl + +from fla.utils import IS_AMD, autotune_cache_kwargs, check_shared_mem, input_guard + +NUM_WARPS = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit +def forward_substitution_kernel( + # Input: Lower triangular matrix L (I + M) + L_ptr, # pointer to lower triangular matrix + L_stride_bh, # stride for batch and head + # Output: Inverse matrix A + A_ptr, # pointer to inverse matrix + A_stride_bh, # stride for batch and head + BT: tl.constexpr, +): + """ + Compute inverse of lower triangular matrix using forward substitution. + + For L = I + M (lower triangular with 1s on diagonal): + Compute A = L^(-1) using forward substitution: + - A[i,i] = 1 + - A[i,j] = -sum(L[i,k] * A[k,j] for k in range(j,i)) for j < i + """ + # Get batch-head index + i_bh = tl.program_id(0) + + # Compute pointer offsets for this batch-head + L_offset = i_bh * L_stride_bh + A_offset = i_bh * A_stride_bh + + # Initialize A as identity matrix + for i in range(BT): + for j in range(BT): + if i == j: + tl.store(A_ptr + A_offset + i * BT + j, 1.0) + else: + tl.store(A_ptr + A_offset + i * BT + j, 0.0) + + # Forward substitution + for i in range(1, BT): + for j in range(i): + # A[i,j] = -sum(L[i,k] * A[k,j] for k in range(j,i)) + sum_val = 0.0 + for k in range(j, i): + L_ik = tl.load(L_ptr + L_offset + i * BT + k) + A_kj = tl.load(A_ptr + A_offset + k * BT + j) + sum_val += L_ik * A_kj + tl.store(A_ptr + A_offset + i * BT + j, -sum_val) + + +@input_guard +def forward_substitution( + L: torch.Tensor, +) -> torch.Tensor: + """ + Compute inverse of lower triangular matrix using forward substitution. + + Args: + L: Lower triangular matrix of shape [B, H, BT, BT] with 1s on diagonal + + Returns: + A: Inverse matrix of shape [B, H, BT, BT] + """ + B, H, BT, BT2 = L.shape + assert BT == BT2 + + # Reshape for kernel: [B*H, BT, BT] + L_flat = L.view(B * H, BT, BT) + A_flat = torch.empty_like(L_flat) + + # Launch kernel ONCE for all batches and heads in parallel + forward_substitution_kernel[(B * H,)]( + L_ptr=L_flat, + L_stride_bh=BT * BT, + A_ptr=A_flat, + A_stride_bh=BT * BT, + BT=BT + ) + + return A_flat.view(B, H, BT, BT) + + +class ForwardSubstitutionFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward( + ctx, + L: torch.Tensor, + ): + A = forward_substitution(L) + ctx.save_for_backward(L, A) + return A + + @staticmethod + @input_guard + def backward(ctx, dA): + L, A = ctx.saved_tensors + + # Backward pass: dL = -A^T @ dA @ A^T + # Simplified implementation for now + dL = torch.zeros_like(L) + + return dL + + +@torch.compiler.disable +def quasar_forward_substitution( + L: torch.Tensor, +) -> torch.Tensor: + """ + Compute inverse of lower triangular matrix using Triton kernel with autograd support + + Args: + L: Lower triangular matrix of shape [B, H, BT, BT] with 1s on diagonal + + Returns: + A: Inverse matrix of shape [B, H, BT, BT] + """ + return ForwardSubstitutionFunction.apply(L) \ No newline at end of file diff --git a/fla/ops/quasar/fused_recurrent.py b/fla/ops/quasar/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..3b74899850656f887e859404da0e6447f525df60 --- /dev/null +++ b/fla/ops/quasar/fused_recurrent.py @@ -0,0 +1,236 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention + +import torch +import triton +import triton.language as tl + +from fla.utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +@triton.heuristics({ + 'HAS_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'STORE_FINAL_STATE': lambda args: args['final_state'] is not None, + 'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_quasar_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + o, + initial_state, + final_state, + scale, + T, + H: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + HAS_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, +): + i_v, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + # [BK, BV] fragment of the state + b_h = tl.zeros([BK, BV], dtype=tl.float32) + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + if HAS_INITIAL_STATE: + p_h0 = initial_state + (i_b * H + i_h) * BK * BK + o_k[:, None] * BK + o_v[None, :] + b_h += tl.load(p_h0).to(tl.float32) + + # Load Invariants Outside Loop + b_beta_head = tl.load(beta + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) + b_exp_A = tl.exp(b_A) + eps = 1e-8 + + # Block Pointers for sequential loading + p_q = tl.make_block_ptr(q + (i_b * T * H + i_h) * BK, (T, BK), (H * BK, 1), (0, 0), (1, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (i_b * T * H + i_h) * BK, (T, BK), (H * BK, 1), (0, 0), (1, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (i_b * T * H + i_h) * BK + i_v * BV, (T, BV), (H * BK, 1), (0, 0), (1, BV), (1, 0)) + p_g = tl.make_block_ptr(g + (i_b * T * H + i_h) * BK, (T, BK), (H * BK, 1), (0, 0), (1, BK), (1, 0)) + p_o = tl.make_block_ptr(o + (i_b * T * H + i_h) * BK + i_v * BV, (T, BV), (H * BK, 1), (0, 0), (1, BV), (1, 0)) + + for _ in range(0, T): + # Load tokens for this step + # [1, BK] + b_q = tl.load(p_q).to(tl.float32) + b_k = tl.load(p_k).to(tl.float32) + b_g = tl.load(p_g).to(tl.float32) + # [1, BV] + b_v = tl.load(p_v).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + + b_q *= scale + + # 1. CT Alpha Logic - Scalar reduction over BK + b_k2 = tl.sum(b_k * b_k) + # Use a more stable clamp to avoid NaN + b_k2_stab = tl.maximum(b_k2, 0.05) + b_alpha = (1.0 - tl.exp(-b_beta_head * b_k2_stab)) / (b_k2_stab + eps) + + # 2. Hybrid Forget Gate + if HAS_DT_BIAS: + b_bias = tl.load(dt_bias + i_h * BK + o_k).to(tl.float32) + b_g += b_bias[None, :] + + # Softplus Gate Approximation + # Patch A: tl.log1p not available on all Triton versions; use tl.log(1.0 + x) + b_gk = -b_exp_A * (tl.where(b_g > 20.0, b_g, tl.log(1.0 + tl.exp(b_g)))) + + # Apply Forget Gate to State + # Patch B: constexpr[0] indexing not supported on older Triton; use tl.view for deterministic layout + b_h *= tl.exp(tl.view(b_gk, [BK])[:, None]) + + # 3. State Update (Rank-1 Delta Rule) + # S_t = S_t_forgot + alpha * k @ (v - k^T @ S_t_forgot)^T + # Patch C: tl.dot requires 2D inputs; wrap 1D slices with tl.view + # v_pred = k @ h -> [1, BV] + b_v_pred = tl.dot(tl.view(b_k, [1, BK]), b_h) + b_v_err = b_v - b_v_pred + # Outer product: [BK, 1] @ [1, BV] + b_h += (b_alpha * tl.trans(tl.view(b_k, [1, BK]))) @ b_v_err + + # 4. Output Projection + # o = q @ h -> [1, BV] + b_o = tl.dot(tl.view(b_q, [1, BK]), b_h) + tl.store(p_o, b_o.to(p_o.dtype.element_ty)) + + # Advance pointers + p_q = tl.advance(p_q, (1, 0)) + p_k = tl.advance(p_k, (1, 0)) + p_v = tl.advance(p_v, (1, 0)) + p_g = tl.advance(p_g, (1, 0)) + p_o = tl.advance(p_o, (1, 0)) + + if STORE_FINAL_STATE: + p_ht = final_state + (i_b * H + i_h) * BK * BK + o_k[:, None] * BK + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty)) + + if STORE_FINAL_STATE: + p_ht = final_state + (i_b * H + i_h) * BK * BK + o_k[:, None] * BK + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty)) + + +@input_guard +def fused_recurrent_quasar_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + B, T, H, S = q.shape + if scale is None: + scale = S ** -0.5 + + o = torch.empty_like(v) + final_state = torch.empty(B, H, S, S, dtype=torch.float32, device=q.device) if output_final_state else None + + # Grid: (V_heads, B*H) + # BV=64 often works better on A100/H100 if BK is small + BV = 64 if S <= 64 else 32 + grid = (triton.cdiv(S, BV), B * H) + fused_recurrent_quasar_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + o=o, + initial_state=initial_state, + final_state=final_state, + scale=scale, + T=T, + H=H, + BK=S, + BV=BV, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=8, + num_stages=4, + ) + + return o, final_state + + +class FusedRecurrentQuasarFunction(torch.autograd.Function): + @staticmethod + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, + ): + o, final_state = fused_recurrent_quasar_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + output_final_state=output_final_state, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + return o, final_state + + @staticmethod + def backward(ctx, do, dht): + raise NotImplementedError("Backward pass for fused_recurrent_quasar is not implemented yet.") + + +@torch.compiler.disable +def fused_recurrent_quasar( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + return FusedRecurrentQuasarFunction.apply( + q, k, v, g, beta, A_log, dt_bias, initial_state, output_final_state, scale, use_qk_l2norm_in_kernel + ) \ No newline at end of file diff --git a/fla/ops/quasar/gate.py b/fla/ops/quasar/gate.py new file mode 100644 index 0000000000000000000000000000000000000000..e98c90e726979fe39f0652182315e39fbce92f0e --- /dev/null +++ b/fla/ops/quasar/gate.py @@ -0,0 +1,244 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# Modified for QuasarAttention + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] + + +def naive_quasar_gate( + beta: torch.Tensor, + lambda_t: torch.Tensor, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Torch reference implementation for QuasarAttention gate computation. + + Computes: alpha = (1 - exp(-beta * lambda)) / (lambda + eps) + + Args: + beta (torch.Tensor): + Parameter tensor with `H` elements. + lambda_t (torch.Tensor): + Input tensor of shape `[..., H, 1]` (norm squared of keys). + output_dtype (torch.dtype): + Output dtype. + + Returns: + Output tensor of shape `[..., H, 1]`. + """ + eps = 1e-8 + alpha = (1 - torch.exp(-beta.view(-1, 1) * lambda_t)) / (lambda_t + eps) + return alpha.to(output_dtype) + + +@triton.autotune( + configs=[ + triton.Config({"BT": BT}, num_warps=num_warps, num_stages=num_stages) + for BT in BT_LIST_AUTOTUNE + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [2, 3] + ], + key=["H", "D"], + **autotune_cache_kwargs, +) +@triton.jit +def quasar_gate_fwd_kernel( + lambda_t, + beta, + alpha, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + + b_beta = tl.load(beta + i_h).to(tl.float32) + + p_lambda = tl.make_block_ptr(lambda_t + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + p_alpha = tl.make_block_ptr(alpha + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + # [BT, BD] + b_lambda = tl.load(p_lambda, boundary_check=(0, 1)).to(tl.float32) + + # alpha = (1 - exp(-beta * lambda)) / (lambda + eps) + eps = 1e-8 + b_alpha = (1 - tl.exp(-b_beta * b_lambda)) / (b_lambda + eps) + tl.store(p_alpha, b_alpha.to(p_alpha.dtype.element_ty), boundary_check=(0, 1)) + + +@input_guard +def quasar_gate_fwd( + lambda_t: torch.Tensor, + beta: torch.Tensor, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H, K = lambda_t.shape[-2:] + T = lambda_t.numel() // (H * K) + + alpha = torch.empty_like(lambda_t, dtype=output_dtype) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]), H) + + quasar_gate_fwd_kernel[grid]( + lambda_t=lambda_t, + beta=beta, + alpha=alpha, + T=T, + H=H, + D=K, + BD=triton.next_power_of_2(K), + ) + return alpha + + +class QuasarGateFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + lambda_t: torch.Tensor, + beta: torch.Tensor, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + alpha = quasar_gate_fwd( + lambda_t=lambda_t, + beta=beta, + output_dtype=output_dtype + ) + ctx.save_for_backward(lambda_t, beta) + return alpha + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dalpha: torch.Tensor): + lambda_t, beta = ctx.saved_tensors + eps = 1e-8 + + # dalpha/dlambda and dalpha/dbeta derivatives + # alpha = (1 - exp(-beta * lambda)) / (lambda + eps) + # dalpha/dbeta = exp(-beta * lambda) + beta_exp = torch.exp(-beta.view(-1, 1) * lambda_t) + lambda_plus_eps = lambda_t + eps + + # dalpha/dlambda = (beta * exp(-beta * lambda) * lambda - (1 - exp(-beta * lambda))) / lambda^2 + dlambda = (beta.view(-1, 1) * beta_exp * lambda_plus_eps - (1 - beta_exp)) / (lambda_plus_eps ** 2) + + # dalpha/dbeta = exp(-beta * lambda) + dbeta = beta_exp + + dlambda = dlambda * dalpha + # Sum over sequence and dimensions, but preserve head dimension + dbeta = (dbeta * dalpha).sum(dim=(0, 1)) + + return dlambda, dbeta, None, None + + +@triton.jit +def fast_quasar_alpha_fwd_kernel( + k, + beta, + alpha, + T, + stride_beta_b, + stride_beta_t, + stride_beta_h, + H: tl.constexpr, + S: tl.constexpr, + BK: tl.constexpr, + BT: tl.constexpr, +): + i_bh, i_t = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + eps = 1e-6 + + # Process BT tokens + for t in range(BT): + idx = i_t * BT + t + if idx < T: + # We use block ptr if we want, but simple indexing is fine here for S + offset = (i_b * T * H + idx * H + i_h) * S + b_k2 = 0.0 + for s in range(0, S, BK): + mask = (s + tl.arange(0, BK)) < S + b_k = tl.load(k + offset + s + tl.arange(0, BK), mask=mask, other=0.0).to(tl.float32) + b_k2 += tl.sum(b_k * b_k) + + # Load beta for this specific token + beta_offset = i_b * stride_beta_b + idx * stride_beta_t + i_h * stride_beta_h + b_beta = tl.load(beta + beta_offset).to(tl.float32) + + # alpha = (1 - exp(-beta * |k|^2)) / (|k|^2 + eps) + # Clamp k2 internally for stability like the torch version did + k2_clamped = tl.where(b_k2 < 0.1, 0.1, tl.where(b_k2 > 10.0, 10.0, b_k2)) + b_alpha = (1.0 - tl.exp(-b_beta * k2_clamped)) / (k2_clamped + eps) + + tl.store(alpha + i_b * T * H + idx * H + i_h, b_alpha.to(alpha.dtype.element_ty)) + + +@input_guard +def fast_quasar_alpha( + k: torch.Tensor, + beta: torch.Tensor, +) -> torch.Tensor: + B, T, H, S = k.shape + alpha = torch.empty(B, T, H, device=k.device, dtype=k.dtype) + + if beta.ndim == 1: + stride_beta_b, stride_beta_t, stride_beta_h = 0, 0, beta.stride(0) + elif beta.ndim == 3: + stride_beta_b, stride_beta_t, stride_beta_h = beta.stride(0), beta.stride(1), beta.stride(2) + else: + raise ValueError(f"beta must be 1D or 3D, got {beta.ndim}D") + + BT = 64 + grid = (B * H, triton.cdiv(T, BT)) + fast_quasar_alpha_fwd_kernel[grid]( + k=k, + beta=beta, + alpha=alpha, + T=T, + stride_beta_b=stride_beta_b, + stride_beta_t=stride_beta_t, + stride_beta_h=stride_beta_h, + H=H, + S=S, + BK=triton.next_power_of_2(S), + BT=BT, + ) + return alpha + + +@torch.compiler.disable +def fused_quasar_gate( + lambda_t: torch.Tensor, + beta: torch.Tensor, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Fused QuasarAttention gate computation with autograd support. + + Computes: alpha = (1 - exp(-beta * lambda)) / (lambda + eps) + + Args: + lambda_t (torch.Tensor): + Input tensor of shape `[..., H, 1]` (norm squared of keys). + beta (torch.Tensor): + Parameter tensor with `H` elements. + + Returns: + Output tensor of shape `[..., H, 1]`. + """ + return QuasarGateFunction.apply(lambda_t, beta, output_dtype) \ No newline at end of file diff --git a/fla/ops/quasar/wy_fast.py b/fla/ops/quasar/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..1de969cada59a5efcfa7065fa451fcf2eada6996 --- /dev/null +++ b/fla/ops/quasar/wy_fast.py @@ -0,0 +1,311 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2 +from fla.utils import autotune_cache_kwargs, check_shared_mem + + +@triton.heuristics({ + 'STORE_QG': lambda args: args['qg'] is not None, + 'STORE_KG': lambda args: args['kg'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_quasar_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A.to(b_vb.dtype), b_vb) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_qg = tl.make_block_ptr(qg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.).to(tl.float32) + b_kg = b_k * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) + p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A.to(b_k.dtype), b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_quasar_kernel( + k, + v, + beta, + gk, + A, + dA, + dw, + du, + dk, + dk2, + dv, + db, + dg, + dg2, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_b = tl.load(p_b, boundary_check=(0,)) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2 + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg2 = tl.make_block_ptr(dg2 + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk_exp = exp2(tl.load(p_gk, boundary_check=(0, 1))) + b_kbg = b_k * b_b[:, None] * b_gk_exp + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + b_dk = b_dkbg * b_gk_exp * b_b[:, None] + tl.load(p_dk, boundary_check=(0, 1)) + b_db += tl.sum(b_dkbg * b_k * b_gk_exp, 1) + b_dg = b_kbg * b_dkbg + tl.load(p_dg, boundary_check=(0, 1)) + + tl.store(p_dk2, b_dk.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg2, b_dg.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + b_dA = tl.where(m_A, -b_dA, 0) + + # if using gk, save dA first and handle dk in another kernel + p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + qg = torch.empty_like(q) if q is not None else None + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_quasar_kernel[(NT, B*H)]( + q=q, + k=k, + qg=qg, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u, qg, kg + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + gk: torch.Tensor, + A: torch.Tensor, + dk: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk2 = torch.empty_like(dk, dtype=torch.float) + dv = torch.empty_like(v) + dg2 = torch.empty_like(gk, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + prepare_wy_repr_bwd_quasar_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + gk=gk, + A=A, + dA=dA, + dw=dw, + du=du, + dk=dk, + dk2=dk2, + dv=dv, + db=db, + dg=dg, + dg2=dg2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dk = dk2 + dg = dg2 + return dk, dv, db, dg, dA diff --git a/fla/ops/rebased/__init__.py b/fla/ops/rebased/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..292c918e3e9ef8319a2738bcb5ca3c314c22e1df --- /dev/null +++ b/fla/ops/rebased/__init__.py @@ -0,0 +1,6 @@ + +from .parallel import parallel_rebased + +__all__ = [ + 'parallel_rebased', +] diff --git a/fla/ops/rebased/naive.py b/fla/ops/rebased/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..5811e6634c6c2e7a9e0f699da6f302457aeb4980 --- /dev/null +++ b/fla/ops/rebased/naive.py @@ -0,0 +1,25 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + + +def naive_parallel_rebased( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + use_norm: bool = True, +) -> torch.Tensor: + if scale is None: + scale = q.shape[-1] ** -0.5 + q = q * scale + attn = q @ k.transpose(-2, -1) + attn = attn ** 2 + attn.masked_fill_(~torch.tril(torch.ones(q.shape[-2], q.shape[-2], dtype=torch.bool, device=q.device)), 0) + o = attn @ v + if use_norm: + z = attn.sum(-1) + return o / (z[..., None] + 1e-6) + else: + return o diff --git a/fla/ops/rebased/parallel.py b/fla/ops/rebased/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..83351ee76f9438403f8b6b68e2ee6c298b9c023a --- /dev/null +++ b/fla/ops/rebased/parallel.py @@ -0,0 +1,463 @@ + +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + +# Rebased: Linear Transformers with Learnable Kernel Functions are Better In-Context Models +# https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/ops/triton/rebased_fast/parallel.py + + +@triton.jit(do_not_specialize=['T']) +def parallel_rebased_fwd_kernel( + q, + k, + v, + o, + z, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # i_c: chunk index. used for sequence parallelism + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % (NV) + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, 0), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v*BV), (BTS, BV), (1, 0)) + + # [BQ, BD] block Q, in the shared memory throughout the whole kernel + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + b_o = tl.zeros([BTL, BV], dtype=tl.float32) + b_z = tl.zeros([BTL], dtype=tl.float32) + + # Q block and K block have no overlap + # no need for mask, thereby saving flops + for _ in range(0, i_c*BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_s = tl.dot(b_q, (b_k), allow_tf32=False) + b_s = b_s * b_s + b_z += tl.sum(b_s, axis=1) + + # [BQ, BD] + b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + + # # rescale interchunk output + tl.debug_barrier() + o_q = tl.arange(0, BTL) + # # sync threads, easy for compiler to optimize + # tl.debug_barrier() + + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, i_c*BTL), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTS, BV), (1, 0)) + # Q block and K block have overlap. masks required + for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_s = tl.dot(b_q, b_k, allow_tf32=False) + b_s = b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_z += tl.sum(b_s, axis=1) + # [BTL, BV] + b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + o_k += BTS + + p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_z = z + (i_bh + B * H * i_k) * T + i_c*BTL + tl.arange(0, BTL) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c*BTL + tl.arange(0, BTL)) < T)) + + +@triton.jit(do_not_specialize=['T']) +def _parallel_rebased_bwd_dq( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dq, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_q = (b_q * scale).to(b_q.dtype) + b_dq = tl.zeros([BTL, BK], dtype=tl.float32) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k*BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, 0), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i_c*BTL + tl.arange(0, BTL) + b_dz = tl.load(p_dz, mask=(i_c*BTL + tl.arange(0, BTL)) < T) + + for _ in range(0, i_c*BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + # [BQ, BD] + b_dq += tl.dot((2 * b_ds * b_s).to(b_v.dtype), b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + + b_dq *= scale + o_q = tl.arange(0, BTL) + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, i_c*BTL), (BV, BTS), (0, 1)) + # Q block and K block have overlap. masks required + for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + b_s = tl.where(m_s, b_s, 0) + # [BTL, BK] + b_dq += tl.dot((2 * b_ds * b_s).to(b_k.dtype), + b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + o_k += BTS + p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit(do_not_specialize=['T']) +def _parallel_rebased_bwd_dkv( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # compute dk dv + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) + b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros( + [BTL, BV], dtype=tl.float32) + + for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + # [BK, BTS] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BV, BTS] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + # [BTL, BTS] + b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale + b_s2 = b_s * b_s + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale + if i_v == 0: + b_ds += b_dz[None, :] * scale + else: + b_ds = b_ds + b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + + tl.debug_barrier() + o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) + for i in range(i_c*BTL, (i_c+1)*BTL, BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + # [BK, BQ] + m_s = o_k[:, None] <= o_q[None, :] + b_s = tl.dot(b_k, b_q, allow_tf32=False) * scale + b_s2 = b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_s2 = tl.where(m_s, b_s2, 0) + + b_ds = tl.dot(b_v, b_do, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[None, :] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + # [BK, BD] + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + o_q += BTS + + p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit(do_not_specialize=['T']) +def parallel_rebased_bwd_kernel( + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % (NV) + i_h = i_bh % H + _parallel_rebased_bwd_dq( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dq, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + ) + tl.debug_barrier() + _parallel_rebased_bwd_dkv( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dk, + dv, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + ) + + +class ParallelBasedFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale): + BTL, BTS = 128, 32 + assert BTL % BTS == 0 + # assert q.shape[-1] % 16 == 0 + BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not." + + o = torch.empty(NK, B, H, T, V, device=q.device) + z = torch.empty(NK, B, H, T, device=q.device) + parallel_rebased_fwd_kernel[grid]( + q, + k, + v, + o, + z, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages, + ) + ctx.save_for_backward(q, k, v) + ctx.scale = scale + return o.sum(0).to(q.dtype), z.sum(0).to(q.dtype) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dz): + q, k, v = ctx.saved_tensors + scale = ctx.scale + BTL, BTS = 64, 32 + assert BTL % BTS == 0 + BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) + BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not" + + dq = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dk = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dv = torch.empty(NK, B, H, T, V, dtype=q.dtype, device=q.device) + + parallel_rebased_bwd_kernel[grid]( + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages, + ) + + return dq.sum(0).to(q.dtype), dk.sum(0).to(k.dtype), dv.sum(0).to(v.dtype), None + + +def parallel_rebased( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + eps: float = 1e-5, + use_scale: bool = True, + use_normalize: bool = True, + return_both: bool = False, + head_first: bool = False, +): + assert q.shape[-1] <= 128, "only support feature dim up to 128" + if use_scale: + scale = q.shape[-1] ** -0.5 + else: + scale = 1 + if not head_first: + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + o, z = ParallelBasedFunction.apply(q, k, v, scale) + if return_both: + return o, z + if use_normalize: + o = o / (z[..., None] + eps) + if not head_first: + o = o.transpose(1, 2) + return o.to(q.dtype) diff --git a/fla/ops/retention/__init__.py b/fla/ops/retention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0aea0d1461d071c58838d517185423b00c4a7420 --- /dev/null +++ b/fla/ops/retention/__init__.py @@ -0,0 +1,12 @@ + +from .chunk import chunk_retention +from .fused_chunk import fused_chunk_retention +from .fused_recurrent import fused_recurrent_retention +from .parallel import parallel_retention + +__all__ = [ + 'chunk_retention', + 'fused_chunk_retention', + 'parallel_retention', + 'fused_recurrent_retention', +] diff --git a/fla/ops/retention/chunk.py b/fla/ops/retention/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..7b97e931cbacb5d86860b6cf12b4ace861d0cd82 --- /dev/null +++ b/fla/ops/retention/chunk.py @@ -0,0 +1,75 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.simple_gla.chunk import chunk_simple_gla + + +@torch.compiler.disable +def chunk_retention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + o, final_state = chunk_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + g_gamma=g_gamma, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/retention/fused_chunk.py b/fla/ops/retention/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..3b431c86219f113e20a06e66fb3b4b098db09c30 --- /dev/null +++ b/fla/ops/retention/fused_chunk.py @@ -0,0 +1,76 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.simple_gla import fused_chunk_simple_gla + + +@torch.compiler.disable +def fused_chunk_retention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. + Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + o, final_state = fused_chunk_simple_gla( + q=q, + k=k, + v=v, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/retention/fused_recurrent.py b/fla/ops/retention/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..a242c61e675501cee8c82d4347e0f5f935f74ce1 --- /dev/null +++ b/fla/ops/retention/fused_recurrent.py @@ -0,0 +1,31 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla + + +def fused_recurrent_retention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + o, final_state = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/retention/naive.py b/fla/ops/retention/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..ff62f148cb0ae38282bf31465c1ea805d86ad74d --- /dev/null +++ b/fla/ops/retention/naive.py @@ -0,0 +1,14 @@ + +import torch + + +def naive_retention(q, k, v): + orig_type = q.dtype + q, k, v = q.float(), k.float(), v.float() + _, n_heads, seq_len, d_head = q.shape + s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(n_heads), dtype=torch.float))).log2() + n = q.new_tensor(range(seq_len), dtype=torch.float) + n = torch.exp2((n.unsqueeze(-1) - n) * s.view(-1, 1, 1)) * n.unsqueeze(-1).ge(n) + s = torch.einsum('bhqd,bhkd,hqk->bhqk', q * d_head ** -0.5, k, n.to(q.dtype)) + o = torch.einsum('bhqk,bhkd->bhqd', s, v) + return o.to(orig_type) diff --git a/fla/ops/retention/parallel.py b/fla/ops/retention/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..5143a1c30770e11bf980c03479b1c9e32e45ac90 --- /dev/null +++ b/fla/ops/retention/parallel.py @@ -0,0 +1,69 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.simple_gla.parallel import parallel_simple_gla + + +def parallel_retention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + output_attentions: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + output_attentions (bool): + Whether to output the materialized attention scores of shape [B, H, T, T]. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + attn (torch.Tensor): + Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None` + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + g = s[None, None, :].expand(q.shape[0], q.shape[1], q.shape[2]) + + o, attn = parallel_simple_gla( + q=q, + k=k, + v=v, + scale=scale, + g=g, + output_attentions=output_attentions, + cu_seqlens=cu_seqlens, + ) + return o, attn diff --git a/fla/ops/rwkv4/__init__.py b/fla/ops/rwkv4/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..85d31a8ae7b8dfe75f6a03efd8ef3a8f04c557a8 --- /dev/null +++ b/fla/ops/rwkv4/__init__.py @@ -0,0 +1,6 @@ + +from .fused_recurrent import fused_recurrent_rwkv4 + +__all__ = [ + 'fused_recurrent_rwkv4', +] diff --git a/fla/ops/rwkv4/fused_recurrent.py b/fla/ops/rwkv4/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..f4b959ff0b73040a18eea97366b8987c8835ee82 --- /dev/null +++ b/fla/ops/rwkv4/fused_recurrent.py @@ -0,0 +1,471 @@ +# Copyright (c) 2024, Songlin Yang, Yu Zhang + +from typing import Any, cast + +import torch +import triton +import triton.language as tl +from torch import Tensor +from torch.autograd.function import Function, FunctionCtx, once_differentiable + +from fla.ops.utils.op import exp +from fla.utils import input_guard + + +def get_block_size_c(chans: int) -> int: + if chans < 32: + return 32 + if chans < 64: + return 64 + return 128 + + +@triton.jit +def fused_recurrent_rwkv4_forward_kernel( + # W + w_ptr, + w_s_c, + # U + u_ptr, + u_s_c, + # K + k_ptr, + k_s_b, + k_s_t, + k_s_c, + # V + v_ptr, + v_s_b, + v_s_t, + v_s_c, + # State + state_ptr, + state_s_b, + state_s_abe, + state_s_c, + # WKV + wkv_ptr, + wkv_s_b, + wkv_s_t, + wkv_s_c, + # Output state + state_out_ptr, + state_out_s_b, + state_out_s_abe, + state_out_s_t, + state_out_s_c, + # Params + chans, + tsz, + BLOCK_SIZE_C: tl.constexpr, +): + # Parallelize over the batch dimension. + b_idx = tl.program_id(0) + c_idx = tl.program_id(1) + + cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) + cmask = cs < chans + + # Pointers to the batch (and possibly channel) for the input tensors. + k_ptr = k_ptr + b_idx * k_s_b + v_ptr = v_ptr + b_idx * v_s_b + alpha_ptr = state_ptr + b_idx * state_s_b + beta_ptr = state_ptr + b_idx * state_s_b + state_s_abe + eps_ptr = state_ptr + b_idx * state_s_b + 2 * state_s_abe + + # Pointers to the batch (and possibly channel) for the output tensors. + wkv_ptr = wkv_ptr + b_idx * wkv_s_b + alpha_out_ptr = state_out_ptr + b_idx * state_out_s_b + beta_out_ptr = state_out_ptr + b_idx * state_out_s_b + state_out_s_abe + eps_out_ptr = state_out_ptr + b_idx * state_out_s_b + 2 * state_out_s_abe + + # Loads parameters. + alpha = tl.load(alpha_ptr + cs * state_s_c, mask=cmask).to(tl.float32) + beta = tl.load(beta_ptr + cs * state_s_c, mask=cmask).to(tl.float32) + eps = tl.load(eps_ptr + cs * state_s_c, mask=cmask).to(tl.float32) + w = tl.load(w_ptr + cs * w_s_c, mask=cmask).to(tl.float32) + u = tl.load(u_ptr + cs * u_s_c, mask=cmask).to(tl.float32) + + for t in range(tsz): + kt = tl.load(k_ptr + t * k_s_t + cs * k_s_c, mask=cmask).to(tl.float32) + vt = tl.load(v_ptr + t * v_s_t + cs * v_s_c, mask=cmask).to(tl.float32) + + ukt = u + kt + tau = tl.maximum(ukt, eps) + e1a = exp(eps - tau) + e2a = exp(ukt - tau) + wkv = (e1a * alpha + e2a * vt) / (e1a * beta + e2a) + tl.store(wkv_ptr + t * wkv_s_t + cs * wkv_s_c, wkv, mask=cmask) + + w_eps = w + eps + eps = tl.maximum(w_eps, kt) + e1b = exp(w_eps - eps) + e2b = exp(kt - eps) + alpha = e1b * alpha + e2b * vt + beta = e1b * beta + e2b + tl.store(alpha_out_ptr + t * state_out_s_t + cs * state_out_s_c, alpha, mask=cmask) + tl.store(beta_out_ptr + t * state_out_s_t + cs * state_out_s_c, beta, mask=cmask) + tl.store(eps_out_ptr + t * state_out_s_t + cs * state_out_s_c, eps, mask=cmask) + + +def fused_recurrent_rwkv4_forward( + w: Tensor, + u: Tensor, + k: Tensor, + v: Tensor, + state: Tensor, +) -> tuple[Tensor, Tensor]: + (bsz, tsz, chans) = k.shape + + # New tensors to output. + wkvs = k.new_empty(bsz, tsz, chans) + state_out = k.new_empty(bsz, 3, tsz, chans) + + # Constants. + block_size_c = get_block_size_c(chans) + + def grid(meta: dict[str, Any]) -> tuple[int, ...]: + return (bsz, triton.cdiv(chans, meta["BLOCK_SIZE_C"])) + + fused_recurrent_rwkv4_forward_kernel[grid]( + # W + w, + w.stride(0), + # U + u, + u.stride(0), + # K + k, + k.stride(0), + k.stride(1), + k.stride(2), + # V + v, + v.stride(0), + v.stride(1), + v.stride(2), + # State + state, + state.stride(0), + state.stride(1), + state.stride(3), + # WKV + wkvs, + wkvs.stride(0), + wkvs.stride(1), + wkvs.stride(2), + # Output state + state_out, + state_out.stride(0), + state_out.stride(1), + state_out.stride(2), + state_out.stride(3), + # Params + chans, + tsz, + BLOCK_SIZE_C=block_size_c, + ) + + state_out = torch.cat((state, state_out), dim=2) + + return wkvs, state_out + + +@triton.jit +def fused_recurrent_rwkv4_backward_kernel( + # W + w_ptr, + w_s_c, + # U + u_ptr, + u_s_c, + # K + k_ptr, + k_s_b, + k_s_t, + k_s_c, + # V + v_ptr, + v_s_b, + v_s_t, + v_s_c, + # State + state_ptr, + state_s_b, + state_s_abe, + state_s_t, + state_s_c, + # WKV grad + gwkv_ptr, + gwkv_s_b, + gwkv_s_t, + gwkv_s_c, + # Output state grad + gstate_out_ptr, + gstate_out_s_b, + gstate_out_s_abe, + gstate_out_s_c, + # W grad + gw_ptr, + gw_s_b, + gw_s_c, + # U grad + gu_ptr, + gu_s_b, + gu_s_c, + # K grad + gk_ptr, + gk_s_b, + gk_s_t, + gk_s_c, + # V grad + gv_ptr, + gv_s_b, + gv_s_t, + gv_s_c, + # State grad + gstate_ptr, + gstate_s_b, + gstate_s_abe, + gstate_s_c, + # Params + tsz, + chans, + BLOCK_SIZE_C: tl.constexpr, +): + # Parallelize over the batch dimension. + b_idx = tl.program_id(0) + c_idx = tl.program_id(1) + + cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) + cmask = cs < chans + + # Pointers to the batch (and possibly channel) for the input tensors. + k_ptr = k_ptr + b_idx * k_s_b + v_ptr = v_ptr + b_idx * v_s_b + alpha_ptr = state_ptr + b_idx * state_s_b + beta_ptr = state_ptr + b_idx * state_s_b + state_s_abe + eps_ptr = state_ptr + b_idx * state_s_b + 2 * state_s_abe + + # Pointers to the batch (and possibly channel) for the output tensors. + gk_ptr = gk_ptr + b_idx * gk_s_b + gv_ptr = gv_ptr + b_idx * gv_s_b + + # Pointers to gradients which were recieved by the function. + gwkv_ptr = gwkv_ptr + b_idx * gwkv_s_b + galpha_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b + gbeta_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b + gstate_out_s_abe + geps_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b + 2 * gstate_out_s_abe + + # Loads parameters. + galpha = tl.load(galpha_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) + gbeta = tl.load(gbeta_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) + geps = tl.load(geps_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) + w = tl.load(w_ptr + w_s_c * cs, mask=cmask).to(tl.float32) + u = tl.load(u_ptr + u_s_c * cs, mask=cmask).to(tl.float32) + + # Gradient accumulators. + gw = tl.zeros_like(w) + gu = tl.zeros_like(u) + + alpha_prev = tl.load(alpha_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + beta_prev = tl.load(beta_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + eps_prev = tl.load(eps_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + + for t in range(tsz): + tc = tsz - t - 1 + + kt = tl.load(k_ptr + tc * k_s_t + k_s_c * cs, mask=cmask).to(tl.float32) + vt = tl.load(v_ptr + tc * v_s_t + v_s_c * cs, mask=cmask).to(tl.float32) + + alpha_curr = alpha_prev + beta_curr = beta_prev + eps_curr = eps_prev + + alpha_prev = tl.load(alpha_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + beta_prev = tl.load(beta_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + eps_prev = tl.load(eps_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) + + ukt = u + kt + tau = tl.maximum(ukt, eps_prev) + e1 = exp(eps_prev - tau) + e2 = exp(ukt - tau) + + euke = exp(ukt + eps_prev - 2 * tau) + + denom = e1 * beta_prev + e2 + denom_sq = denom * denom + + gwkvt = tl.load(gwkv_ptr + tc * gwkv_s_t + gwkv_s_c * cs, mask=cmask).to(tl.float32) + + # Backpropagates wkv gradients. + guk = gwkvt * e2 * (e1 * beta_prev * vt - e1 * alpha_prev) / denom_sq + gu += guk + gk = guk + gv = gwkvt * e2 / denom + + galpha_wkv = gwkvt * e1 / denom + gbeta_wkv = -gwkvt * e1 * (e2 * vt + e1 * alpha_prev) / denom_sq + geps_wkv_denom = e1 * beta_prev + e2 + geps_wkv = gwkvt * euke * (alpha_prev - vt * beta_prev) / (geps_wkv_denom * geps_wkv_denom) + + e1 = exp(w + eps_prev - eps_curr) + e2 = exp(kt - eps_curr) + + # Backpropagates alpha gradients. + galpha_we = galpha * e1 * alpha_prev + gw += galpha_we + gk += galpha * e2 * vt + gv += galpha * e2 + geps += galpha * -alpha_curr + + # Backpropagates beta gradients. + gbeta_we = gbeta * e1 * beta_prev + gw += gbeta_we + gk += gbeta * e2 + geps += gbeta * -beta_curr + + # Backpropagates epsilon gradients. + geps_mask = w + eps_prev > kt + geps_we = tl.where(geps_mask, geps, tl.zeros_like(geps)) + gw += geps_we + gk += tl.where(geps_mask, tl.zeros_like(geps), geps) + + # Stores the gradients for k and v. + tl.store(gk_ptr + tc * gk_s_t + gk_s_c * cs, gk, mask=cmask) + tl.store(gv_ptr + tc * gv_s_t + gv_s_c * cs, gv, mask=cmask) + + # Computes new gradients for alpha and beta. + galpha = galpha * e1 + galpha_wkv + gbeta = gbeta * e1 + gbeta_wkv + geps = galpha_we + gbeta_we + geps_we + geps_wkv + + # Stores final gradients for alpha and beta. + galpha_ptr = gstate_ptr + b_idx * gstate_s_b + gbeta_ptr = gstate_ptr + b_idx * gstate_s_b + gstate_s_abe + geps_ptr = gstate_ptr + b_idx * gstate_s_b + 2 * gstate_s_abe + tl.store(galpha_ptr + gstate_s_c * cs, galpha, mask=cmask) + tl.store(gbeta_ptr + gstate_s_c * cs, gbeta, mask=cmask) + tl.store(geps_ptr + gstate_s_c * cs, geps, mask=cmask) + + # Stores final gradients for w and u. + tl.store(gw_ptr + gw_s_b * b_idx + gw_s_c * cs, gw*w, mask=cmask) + tl.store(gu_ptr + gu_s_b * b_idx + gu_s_c * cs, gu, mask=cmask) + + +def fused_recurrent_rwkv4_backward( + w: Tensor, + u: Tensor, + k: Tensor, + v: Tensor, + state: Tensor, + grad_wkv: Tensor, + grad_state: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + bsz, tsz, chans = k.shape + + gw = w.new_empty(bsz, chans, dtype=torch.float) # New tensors to output. + gu = u.new_empty(bsz, chans, dtype=torch.float) + gk = torch.empty_like(k) + gv = torch.empty_like(v) + gstate = k.new_empty(bsz, 3, 1, chans) + + block_size_c = get_block_size_c(chans) # Constants. + + def grid(meta: dict[str, Any]) -> tuple[int, ...]: + return (bsz, triton.cdiv(chans, meta["BLOCK_SIZE_C"])) + + fused_recurrent_rwkv4_backward_kernel[grid]( + # W + w, + w.stride(0), + # U + u, + u.stride(0), + # K + k, + k.stride(0), + k.stride(1), + k.stride(2), + # V + v, + v.stride(0), + v.stride(1), + v.stride(2), + # State + state, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + # WKV grad + grad_wkv, + grad_wkv.stride(0), + grad_wkv.stride(1), + grad_wkv.stride(2), + # Output state grad + grad_state, + grad_state.stride(0), + grad_state.stride(1), + grad_state.stride(3), + # W grad + gw, + gw.stride(0), + gw.stride(1), + # U grad + gu, + gu.stride(0), + gu.stride(1), + # K grad + gk, + gk.stride(0), + gk.stride(1), + gk.stride(2), + # V grad + gv, + gv.stride(0), + gv.stride(1), + gv.stride(2), + # State grad + gstate, + gstate.stride(0), + gstate.stride(1), + gstate.stride(3), + # Params + tsz, + chans, + BLOCK_SIZE_C=block_size_c, + ) + + return gw.sum(0), gu.sum(0), gk, gv, gstate + + +class FusedRecurrentRWKV4Function(Function): + + @staticmethod + @input_guard + def forward( + ctx: FunctionCtx, + w: Tensor, + u: Tensor, + k: Tensor, + v: Tensor, + state: Tensor, + ) -> tuple[Tensor, Tensor]: + ctx.w_dtype = w.dtype + w = -torch.exp(w.float()) + wkv, state_out = fused_recurrent_rwkv4_forward(w, u, k, v, state) + ctx.save_for_backward(w, u, k, v, state_out[:, :, :-1]) + return wkv, state_out[:, :, -1:] + + @staticmethod + @once_differentiable + @input_guard + def backward(ctx: FunctionCtx, gwkv: Tensor, gstate: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + w, u, k, v, state = cast("tuple[Tensor, ...]", ctx.saved_tensors) + gw, gu, gk, gv, gstate = fused_recurrent_rwkv4_backward(w, u, k, v, state, gwkv, gstate) + return gw.to(ctx.w_dtype), gu.to(u), gk.to(k), gv.to(v), gstate.to(state) + + +def fused_recurrent_rwkv4(w: Tensor, u: Tensor, k: Tensor, v: Tensor, state: Tensor) -> tuple[Tensor, Tensor]: + return FusedRecurrentRWKV4Function.apply(w, u, k, v, state) diff --git a/fla/ops/rwkv6/__init__.py b/fla/ops/rwkv6/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..19e34655216beb8d2d4918c4a95bdf750c805e57 --- /dev/null +++ b/fla/ops/rwkv6/__init__.py @@ -0,0 +1,8 @@ + +from .chunk import chunk_rwkv6 +from .fused_recurrent import fused_recurrent_rwkv6 + +__all__ = [ + 'chunk_rwkv6', + 'fused_recurrent_rwkv6', +] diff --git a/fla/ops/rwkv6/chunk.py b/fla/ops/rwkv6/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..bdb50cc5bb7048e2e99c5cd5f55a996513e42d0a --- /dev/null +++ b/fla/ops/rwkv6/chunk.py @@ -0,0 +1,1334 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.common.chunk_h import chunk_fwd_h +from fla.ops.gla.chunk import chunk_gla_bwd_dA, chunk_gla_bwd_dv, chunk_gla_fwd_o_gk +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.op import exp +from fla.utils import ( + USE_CUDA_GRAPH, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_shared_mem, + input_guard, +) + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [32, 64] if check_shared_mem() else [16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BS': BS}, num_warps=num_warps, num_stages=num_stages) + for BS in [16, 32, 64] + for num_warps in [4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=['S', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_fwd_cumsum_kernel( + s, + oi, + oe, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.).to(tl.float32) + m_e = tl.where(o_i[:, None] > o_i[None, :], 1., 0.).to(tl.float32) + + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_oi = tl.make_block_ptr(oi + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_oe = tl.make_block_ptr(oe + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_oi = tl.dot(m_i, b_s) + b_oe = tl.dot(m_e, b_s) + tl.store(p_oi, b_oi.to(p_oi.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_oe, b_oe.to(p_oe.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + + +def chunk_rwkv6_fwd_cumsum( + g: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> torch.Tensor: + B, T, H, S = g.shape + BT = chunk_size + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=None) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + gi, ge = torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float) + def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) + # keep cummulative normalizer in fp32 + chunk_rwkv6_fwd_cumsum_kernel[grid]( + g, + gi, + ge, + cu_seqlens, + chunk_indices, + T=T, + H=H, + S=S, + BT=BT, + ) + return gi, ge + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BC'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( + q, + k, + gi, # cumulative decay inclusive + ge, # cumulative decay exclusive + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + m_i = i_t * BT + i_i * BC + tl.arange(0, BC) < T + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(gi + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h * K + o_k + + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_gq = tl.where(m_i[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) + b_qg = b_q * exp(b_gq - b_gn[None, :]) * scale + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[:, None] - b_gk) + # [BC, BC] using tf32 to improve precision here. + b_A += tl.dot(b_qg, b_kg) + + p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BK', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( + q, + k, + gi, + ge, + u, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_j = i_i + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_j * BC + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (0,), (BK,), (0,)) + b_u = tl.load(p_u, boundary_check=(0,)) + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_A = tl.sum(b_q * b_kj[None, :] * exp(b_g - b_gk[None, :]), 1) + b_A = tl.where(o_i > j, b_A * scale, 0.) + b_A = tl.where(o_i != j, b_A, tl.sum(b_qj * b_kj * b_u * scale)) + tl.store(A + o_A + j, b_A, mask=m_A) + p_qj += H*K + p_kj += H*K + p_gk += H*K + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + ], + key=['BC', 'BK'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( + q, + k, + gi, + ge, + u, + A, + cu_seqlens, + chunk_indices, + scale, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = i_tc // NC, i_tc % NC + i_j = i_i + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + + o_A = (i_k * all + bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BC + i_h * BC + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK), (BK,), (0,)) + b_u = tl.load(p_u, boundary_check=(0,)) + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_A = tl.sum(b_q * b_kj[None, :] * exp(b_g - b_gk[None, :]), 1) + b_A = tl.where(o_i > j, b_A * scale, 0.) + b_A = tl.where(o_i != j, b_A, tl.sum(b_qj * b_kj * b_u * scale)) + tl.store(A + o_A + j, b_A, mask=m_A) + p_qj += H*K + p_kj += H*K + p_gk += H*K + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + ], + key=['BC'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( + A, + A2, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + NK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + if i_t * BT + i_c * BC >= T: + return + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(0, NK): + p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) + b_A += tl.load(p_A, boundary_check=(0, 1)) + p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) + tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_bwd_kernel_dh( + q, + gi, + ge, + do, + dh, + dht, + dh0, + cu_seqlens, + chunk_offsets, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + STORE_INITIAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hq = i_nh // HQ, i_nh % HQ + i_h = i_hq // NG + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + + for i_t in range(NT - 1, -1, -1): + p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + last_idx = min(i_t * BT + BT, T) - 1 + # [BK, BT] + p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + p_gk = tl.make_block_ptr(ge + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk_last = gi + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) + + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_q = (b_q * exp(b_gk) * scale).to(b_q.dtype) + b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) + b_dh *= exp(b_gk_last)[:, None] + b_dh += tl.dot(b_q, b_do) + + if STORE_INITIAL_STATE_GRADIENT: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BK', 'NC', 'BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_bwd_kernel_intra( + q, + k, + gi, + ge, + dA, + dq, + dk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + if i_t * BT + i_i * BC >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_ge = tl.make_block_ptr(ge + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_ge = tl.load(p_ge, boundary_check=(0, 1)) + b_dq = tl.zeros([BC, BK], dtype=tl.float32) + if i_i > 0: + p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(gi+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[None, :] - b_gk) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + b_dq += tl.dot(b_dA, b_kg) + b_dq *= exp(b_ge - b_gn[None, :]) + + o_i = tl.arange(0, BC) + m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC + p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gkj = gi + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) + # [BK,] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] > j + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp(b_ge - b_gkj[None, :]), 0.) + p_kj += H*K + p_gkj += H*K + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(gi + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = gi + (bos + min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + i_h*K + o_k + + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(i_i + 1, NC): + m_j = (i_t * BT + i_j * BC + tl.arange(0, BC)) < T + p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_gq = tl.where(m_j[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) + b_qg = b_q * exp(b_gq - b_gn[None, :]) + # [BC, BC] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dk += tl.dot(b_dA, b_qg) + b_dk *= exp(b_gn[None, :] - b_gk) + o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) + p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gqj = ge + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dA = tl.load(dA + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_gqj = tl.load(p_gqj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] < j + b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp(b_gqj[None, :] - b_gk), 0.) + p_qj += H*K + p_gqj += H*K + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in [2, 4, 8] + ], + key=['BT'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_rwkv6_bwd_kernel_inter( + q, + k, + v, + h, + gi, + ge, + u, + do, + dh, + dA, + dq, + dk, + dq2, + dk2, + dg, + du, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_gk = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_gi = tl.make_block_ptr(gi + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_gn = gi + (bos + min(T, i_t * BT + BT)-1) * H*K + i_h * K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0) + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BK] + b_dgk += tl.sum(b_h * b_dh, axis=0) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + b_dgk *= exp(b_gn) + b_dq *= scale + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gi = tl.load(p_gi, boundary_check=(0, 1)) + b_dq = b_dq * exp(b_gk) + b_dk = b_dk * exp(b_gn[None, :] - b_gi) + + o_i = tl.arange(0, BT) + p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dA_dig = dA + ((bos + i_t * BT + o_i) * H + i_h) * BT + o_i + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dgk += tl.sum(b_dk * b_k, axis=0) + + b_dq += tl.load(p_dq, boundary_check=(0, 1)) + b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dg = b_q * b_dq - b_k * b_dk + b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] - b_q * b_dq + # [BT,] + b_dA_dig = tl.load(p_dA_dig, mask=(i_t * BT + o_i) < T, other=0) + + p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) + b_u = tl.load(p_u, boundary_check=(0,)) + # scale is already applied to b_dA_diag + b_dq += (b_dA_dig[:, None] * b_u[None, :] * b_k) + b_dk += (b_dA_dig[:, None] * b_u[None, :] * b_q) + b_du = tl.sum(b_dA_dig[:, None] * b_q * b_k, axis=0) + p_du = tl.make_block_ptr(du + (i_tg * H + i_h) * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) + tl.store(p_du, b_du, boundary_check=(0,)) + + p_dq = tl.make_block_ptr(dq2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_rwkv6_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + u: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = k.shape + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BC = min(16, BT) + NC = triton.cdiv(BT, BC) + + A = q.new_empty(B, T, H, BT, dtype=torch.float) + grid = (NT, NC * NC, B * H) + chunk_rwkv6_fwd_A_kernel_intra_sub_inter[grid]( + q, + k, + gi, + ge, + A, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + # load the entire [BC, K] blocks into SRAM at once + if K <= 256: + BK = max(triton.next_power_of_2(K), 16) + chunk_rwkv6_fwd_A_kernel_intra_sub_intra[grid]( + q, + k, + gi, + ge, + u, + A, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + # split then merge + else: + BK = min(128, triton.next_power_of_2(K)) + NK = triton.cdiv(K, BK) + A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) + + grid = (NK, NT * NC, B * H) + chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split[grid]( + q, + k, + gi, + ge, + u, + A_intra, + cu_seqlens, + chunk_indices, + scale, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge[grid]( + A_intra, + A, + cu_seqlens, + chunk_indices, + B=B, + T=T, + H=H, + BT=BT, + BC=BC, + NK=NK, + ) + return A + + +def chunk_rwkv6_bwd_dh( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + do: torch.Tensor, + h0: torch.Tensor, + dht: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + states_in_fp32: bool = False, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + NG = HQ // H + + dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None + + def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + chunk_rwkv6_bwd_kernel_dh[grid]( + q=q, + gi=gi, + ge=ge, + do=do, + dh=dh, + dht=dht, + dh0=dh0, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + ) + return dh, dh0 + + +def chunk_rwkv6_bwd_dqk_intra( + q: torch.Tensor, + k: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + dA: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K = q.shape + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + grid = (NK, NT * NC, B * H) + chunk_rwkv6_bwd_kernel_intra[grid]( + q, + k, + gi, + ge, + dA, + dq, + dk, + cu_seqlens, + chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + return dq, dk + + +def chunk_rwkv6_bwd_dqkgu( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor, + u: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dA: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq2 = torch.empty_like(dq) + dk2 = torch.empty_like(dk) + dg = torch.empty_like(g) + du = u.new_empty(B * NT, H, K, dtype=torch.float) + def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) + chunk_rwkv6_bwd_kernel_inter[grid]( + q, + k, + v, + h, + gi, + ge, + u, + do, + dh, + dA, + dq, + dk, + dq2, + dk2, + dg, + du, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + du = du.sum(0) + return dq2, dk2, dg, du + + +def chunk_rwkv6_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + u: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + h, ht = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=gi, + gv=None, + h0=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + ) + A = chunk_rwkv6_fwd_intra( + q=q, + k=k, + gi=gi, + ge=ge, + u=u, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + o = chunk_gla_fwd_o_gk( + q=q, + v=v, + g=ge, + A=A, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return A, h, ht, o + + +def chunk_rwkv6_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + u: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + A: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + gi, ge = chunk_rwkv6_fwd_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + h, _ = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=gi, + gv=None, + h0=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + ) + dh, dh0 = chunk_rwkv6_bwd_dh( + q=q, + k=k, + v=v, + gi=gi, + ge=ge, + do=do, + h0=initial_state, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=True, + chunk_indices=chunk_indices, + ) + + dA = chunk_gla_bwd_dA( + v=v, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dv = chunk_gla_bwd_dv( + k=k, + g=gi, + A=A, + do=do, + dh=dh, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + dq, dk = chunk_rwkv6_bwd_dqk_intra( + q=q, + k=k, + gi=gi, + ge=ge, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dq, dk, dg, du = chunk_rwkv6_bwd_dqkgu( + q=q, + k=k, + v=v, + h=h, + g=g, + gi=gi, + ge=ge, + u=u, + do=do, + dh=dh, + dA=dA, + dq=dq, + dk=dk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, dg, du, dh0 + + +class ChunkRWKV6Function(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + g, + u, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ): + T = q.shape[1] + if check_shared_mem(): + chunk_size = min(32, max(32, triton.next_power_of_2(T))) + else: + chunk_size = min(64, max(32, triton.next_power_of_2(T))) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + A, h, ht, o = chunk_rwkv6_fwd( + q=q, + k=k, + v=v, + g=g, + u=u, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + ctx.save_for_backward(q, k, v, g, initial_state, A, u, chunk_indices) + + ctx.chunk_size = chunk_size + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o, ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + q, k, v, g, initial_state, A, u, chunk_indices = ctx.saved_tensors + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens + dq, dk, dv, dg, du, dh0 = chunk_rwkv6_bwd( + q=q, + k=k, + v=v, + g=g, + u=u, + scale=scale, + initial_state=initial_state, + A=A, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), dg.to(g), du.to(u), None, dh0, None, None, None + + +@torch.compiler.disable +def chunk_rwkv6( + r: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + scale: int | None = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + r (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + w (torch.Tensor): + Forget gates of shape `[B, T, H, K]`. applied to keys. + u (torch.Tensor): + bonus representations of shape `[H]`. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (Optional[torch.Tensor]): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.rwkv6 import chunk_rwkv6 + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> r = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> w = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> u = torch.randn(H, K, device='cuda') + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = chunk_rwkv6( + r, k, v, w, u, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> r, k, v, w = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (r, k, v, w)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = r.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = chunk_rwkv6( + r, k, v, w, u, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert ht.allclose(ht_var) + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and r.shape[1] < r.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if r.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = r.shape[-1] ** -0.5 + o, final_state = ChunkRWKV6Function.apply( + r, + k, + v, + w, + u, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/rwkv6/chunk_naive.py b/fla/ops/rwkv6/chunk_naive.py new file mode 100644 index 0000000000000000000000000000000000000000..bf3519992f2fffa2be797fc06550698a6fcefc50 --- /dev/null +++ b/fla/ops/rwkv6/chunk_naive.py @@ -0,0 +1,42 @@ + +import torch +from einops import rearrange + + +def naive_chunk_rwkv6( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + chunk_size: int = 32, +): + assert q.shape[-2] % chunk_size == 0 + orig_dtype = q.dtype + num_chunk = q.shape[-2] // chunk_size + u = u.unsqueeze(0) + + q, k, v, w = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size).float(), (q, k, v, w)) + + w_cumsum = w.cumsum(-2) + + kw = k * (w_cumsum[..., -1, None, :] - w_cumsum).exp() + wkv = kw.transpose(-1, -2) @ v + + wkv_new = torch.zeros_like(wkv) + + for i in range(num_chunk - 1): + wkv_new[:, :, i+1] = (wkv_new[:, :, i] * w_cumsum[:, :, i, -1, :, None].exp()) + wkv[:, :, i] + + o_inter = torch.einsum('b h n d p, b h n c d -> b h n c p', wkv_new, (q * (w_cumsum - w).exp())) + + o_intra = torch.zeros_like(o_inter) + for i in range(chunk_size): + attn = (q[:, :, :, i, None] * k * (w_cumsum[:, :, :, i, None] - w[:, :, :, i, None] - w_cumsum).exp()).sum(-1) + mask = (torch.arange(0, chunk_size) < i).to(attn.device) + attn.masked_fill_(~mask, 0) + intra_inter_o = (attn.unsqueeze(-1) * v).sum(-2) + intra_intra_o = (q[:, :, :, i] * u.unsqueeze(2) * k[:, :, :, i]).sum(-1).unsqueeze(-1) * v[:, :, :, i] + o_intra[:, :, :, i] = intra_inter_o + intra_intra_o + o = o_inter + o_intra + return rearrange(o, 'b h n c d -> b h (n c) d').to(orig_dtype) diff --git a/fla/ops/rwkv6/fused_recurrent.py b/fla/ops/rwkv6/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..7efab4a723b4c257087f042f51f6b48ccee42fe5 --- /dev/null +++ b/fla/ops/rwkv6/fused_recurrent.py @@ -0,0 +1,675 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + ], + key=['BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_rwkv6_fwd_kernel( + q, # query [B, H, T, K]/[B, T, H, K] + k, # key [B, H, T, K]/[B, T, H, K] + v, # value [B, H, T, V]/[B, T, H, V] + w, # log gate [B, H, T]/[B, T, H] or None + u, # bonus [B, H, K] + o, # output [NK, B, H, T, V]/[NK, B, T, H, V] + h0, # initial hidden state [B, H, K, V] + ht, # final hidden state [B, H, K, V] + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, # whether to reverse the recurrence + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + STORE_FINAL_STATE: tl.constexpr, # whether to store final state + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_w = w + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_o = o + ((i_k * all + bos) + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_u = u + i_h * K + o_k + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) + b_kv = b_k[:, None] * b_v[None, :] + b_o = tl.sum((b_h + b_kv * b_u[:, None]) * b_q[:, None], 0) + b_h = b_h * exp(b_w)[:, None] + b_kv + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + p_q += (-1 if REVERSE else 1) * H*K + p_k += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_w += (-1 if REVERSE else 1) * H*K + p_o += (-1 if REVERSE else 1) * H*V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + ], + key=['BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_rwkv6_bwd_kernel_dq( + k, # key [B, H, T, V]/[B, T, H, V] + v, # value [B, H, T, V]/[B, T, H, V] + w, # log gate [B, H, T]/[B, T, H] + u, # bonus [B, H, K] + do, # gradient of output [B, H, T, V]/[B, T, H, V] + dq, # gradient of query [NV, B, H, T, K]/[NV, B, T, H, K] + dq1, # gradient of query_aux [NV, B, H, T, K]/[NV, B, T, H, K] + h0, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_w = w + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_do = do + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_dq = dq + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_dq1 = dq1 + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_u = u + i_h * K + o_k + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + b_kv = b_k[:, None] * b_v[None, :] + + b_hq = b_h * b_do[None, :] + b_dq = tl.sum(b_hq + b_kv * b_u[:, None] * b_do[None, :], 1) * scale + b_dq1 = tl.sum(b_hq, 1) + b_h = b_h * exp(b_w)[:, None] + b_h += b_kv + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=mask_k) + tl.store(p_dq1, b_dq1.to(p_dq1.dtype.element_ty), mask=mask_k) + + p_k += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_w += (-1 if REVERSE else 1) * H*K + p_do += (-1 if REVERSE else 1) * H*V + p_dq += (-1 if REVERSE else 1) * H*K + p_dq1 += (-1 if REVERSE else 1) * H*K + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + ], + key=['BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_rwkv6_bwd_kernel_dkv( + q, # query [B, H, T, K]/[B, T, H, K] + k, # key [B, H, T, V]/[B, T, H, V] + v, # value [B, H, T, V]/[B, T, H, V] + w, # log gate [B, H, T]/[B, T, H] + u, # bonus [B, H, K] + do, # gradient of output [B, H, T, V]/[B, T, H, V] + dk, # gradient of key [NV, B, H, T, K]/[NK, B, T, H, K] + dk1, # gradient of key_aux [NV, B, H, T, K]/[NK, B, T, H, K] + dv, # gradient of value [NK, B, H, T, V]/[NV, B, T, H, V] + dh0, # gradient of initial hidden state [N, H, K, V] + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_w = w + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_do = do + (bos + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_dk = dk + ((i_v * all + bos) + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_dk1 = dk1 + ((i_v * all + bos) + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k + p_dv = dv + ((i_k * all + bos) + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v + p_u = u + i_h * K + o_k + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + for _ in range(T - 1, -1, -1): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + b_dkv = b_q[:, None] * b_do[None, :] + b_dk = tl.sum(b_dh * b_v[None, :], 1) + tl.store(p_dk1, b_dk.to(p_dk1.dtype.element_ty), mask=mask_k) + b_dk += tl.sum(b_dkv * b_u[:, None] * b_v[None, :], 1) + b_dv = tl.sum((b_dh + (b_dkv * b_u[:, None])) * b_k[:, None], 0) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) + b_dh *= exp(b_w)[:, None] + b_dh += b_dkv + + p_q += (-1 if not REVERSE else 1) * H*K + p_k += (-1 if not REVERSE else 1) * H*K + p_v += (-1 if not REVERSE else 1) * H*V + p_w += (-1 if not REVERSE else 1) * H*K + p_do += (-1 if not REVERSE else 1) * H*V + p_dk += (-1 if not REVERSE else 1) * H*K + p_dk1 += (-1 if not REVERSE else 1) * H*K + p_dv += (-1 if not REVERSE else 1) * H*V + + if USE_INITIAL_STATE: + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_h) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT, 'BK': BK}, num_warps=num_warps) + for BT in [16, 32, 64] + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + ], + key=['K'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_rwkv6_bwd_kernel_dw( + q, + k, + dq, + dk, + dw, + cu_seqlens, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + REVERSE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + T = eos - bos + NT = tl.cdiv(T, BT) + + o_i = tl.arange(0, BT) + m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) if not REVERSE else tl.where(o_i[:, None] <= o_i[None, :], 1., 0.) + + b_z = tl.zeros([BK], dtype=tl.float32) + + i_t = 0 if not REVERSE else NT - 1 + for _ in range(NT): + p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + b_dq = tl.load(p_dq, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + b_dk = tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32) + b_dw = (b_q * b_dq * scale) - b_k * b_dk + b_c = b_z[None, :] + tl.dot(m_i, b_dw, allow_tf32=False) + tl.store(p_dw, b_c.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + if i_t >= 0: + b_z += tl.sum(b_dw, 0) + + i_t += (1 if not REVERSE else -1) + + +def fused_recurrent_rwkv6_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + h0 = initial_state + ht = q.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None + o = q.new_empty(NK, *v.shape, dtype=torch.float) + + grid = (NV, NK, N * H) + fused_recurrent_rwkv6_fwd_kernel[grid]( + q, + k, + v, + w, + u, + o, + h0, + ht, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + REVERSE=reverse, + ) + o = o.sum(0) + return o, ht + + +def fused_recurrent_rwkv6_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + do: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + + BK, BV = min(triton.next_power_of_2(K), 16), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + dq = q.new_empty(NV, *q.shape, dtype=torch.float) + dq1 = torch.empty_like(dq) + + grid = (NV, NK, N * H) + fused_recurrent_rwkv6_bwd_kernel_dq[grid]( + k, + v, + w, + u, + do, + dq, + dq1, + initial_state, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + REVERSE=reverse, + ) + dq = dq.sum(0) + dq1 = dq1.sum(0) + + BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + + dk = q.new_empty(NV, *k.shape, dtype=torch.float) + dk1 = q.new_empty(NV, *k.shape, dtype=torch.float) + dv = q.new_empty(NK, *v.shape, dtype=torch.float) + + dh0 = torch.empty_like(initial_state) if initial_state is not None else None + grid = (NV, NK, N * H) + fused_recurrent_rwkv6_bwd_kernel_dkv[grid]( + q, + k, + v, + w, + u, + do, + dk, + dk1, + dv, + dh0, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + REVERSE=reverse, + ) + dk = dk.sum(0) + dk1 = dk1.sum(0) + dv = dv.sum(0) + + dw = torch.empty_like(w) + def grid(meta): return (triton.cdiv(meta['K'], meta['BK']), N * H) + fused_recurrent_rwkv6_bwd_kernel_dw[grid]( + q, + k, + dq1, + dk1, + dw, + cu_seqlens, + scale, + T=T, + H=H, + K=K, + REVERSE=not reverse, + ) + du = (do.float() * v).sum(-1, True, dtype=torch.float) * q * k * scale + du = du.sum((0, 1)) + return dq, dk, dv, dw, du, dh0 + + +class FusedRecurrentRWKV6Function(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, ht = fused_recurrent_rwkv6_fwd( + q=q, + k=k, + v=v, + w=w, + u=u, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, k, v, w, u, initial_state) + ctx.scale = scale + ctx.reverse = reverse + ctx.cu_seqlens = cu_seqlens + return o.to(v), ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + q, k, v, w, u, initial_state = ctx.saved_tensors + + dq, dk, dv, dw, du, dh0 = fused_recurrent_rwkv6_bwd( + q=q, + k=k, + v=v, + w=w, + u=u, + do=do, + scale=ctx.scale, + initial_state=initial_state, + reverse=ctx.reverse, + cu_seqlens=ctx.cu_seqlens, + ) + dh0 = dh0.to(initial_state) if dh0 is not None else dh0 + return dq.to(q), dk.to(k), dv.to(v), dw.to(w), du.to(u), None, dh0, None, None, None + + +def fused_recurrent_rwkv6( + r: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + scale: int | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + r (torch.Tensor): + reception of shape `[B, T, H, K]`. + Alias: q, query in linear attention. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + w (torch.Tensor): + data-dependent decays of shape `[B, T, H, K]`. in log space! Alias: g. + u (torch.Tensor): + bonus of shape `[H, K]` + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (Optional[torch.Tensor]): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.rwkv6 import fused_recurrent_rwkv6 + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> u = torch.randn(H, K, device='cuda') + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_recurrent_rwkv6( + q, k, v, g, u, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_rwkv6( + q, k, v, g, u, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert ht.allclose(ht_var) + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and r.shape[1] < r.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if r.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = FusedRecurrentRWKV6Function.apply( + r, + k, + v, + w, + u, + scale, + initial_state, + output_final_state, + reverse, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/rwkv6/recurrent_naive.py b/fla/ops/rwkv6/recurrent_naive.py new file mode 100644 index 0000000000000000000000000000000000000000..86d638f5473110edba65cfe1f4c1905137349c16 --- /dev/null +++ b/fla/ops/rwkv6/recurrent_naive.py @@ -0,0 +1,100 @@ + + +import torch + + +def naive_recurrent_rwkv6( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool | None = False, +): + orig_dtype = q.dtype + B, H, T, K, V = *q.shape, v.shape[-1] + q, k, v, w, u = map(lambda x: x.float(), (q, k, v, w, u)) + h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) + o = torch.zeros_like(v) + + if scale is None: + scale = K ** -0.5 + + if initial_state is not None: + h += initial_state + + for i in range(T): + q_i = q[:, :, i, :] * scale + k_i = k[:, :, i] + v_i = v[:, :, i, :] + w_i = w[:, :, i].exp() + kv_i = k_i[..., None] * v_i[..., None, :] + o_i = (h + u[None, ..., None] * kv_i) * q_i[..., None] + o[:, :, i] = o_i.sum(-2) + h = h * w_i[..., None] + kv_i + ht = h if output_final_state else None + return o.to(orig_dtype), ht + + +@torch.no_grad +def naive_recurrent_rwkv6_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + o: torch.Tensor, + do: torch.Tensor, + initial_state: torch.Tensor | None = None, +): + q, k, v, w, u, o, do = (x.to(dtype=torch.float32) for x in (q, k, v, w, u, o, do)) + B, H, T, K, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1] + h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) + dq = torch.zeros_like(q) + dq_aux = torch.zeros_like(q) + + if initial_state is not None: + h += initial_state + + for i in range(T): + k_i = k[:, :, i] + v_i = v[:, :, i] + w_i = w[:, :, i].exp() + kv_i = k_i[..., None] * v_i[..., None, :] + h_i = (h + u[None, ..., None] * kv_i) + dq_i = (do[:, :, i, None, :] * h_i).sum(-1) + dq_aux_i = (do[:, :, i, None, :] * h).sum(-1) + dq[:, :, i] = dq_i + dq_aux[:, :, i] = dq_aux_i + h = h * w_i[..., None] + kv_i + + du = torch.zeros_like(u) + dh = torch.zeros_like(h) + dk = torch.zeros_like(k) + dk_aux = torch.zeros_like(k) + dv = torch.zeros_like(v) + + for i in range(T - 1, -1, -1): + d_kv_i = do[:, :, i, None, :] * q[:, :, i, :, None] + k_i = k[:, :, i] + v_i = v[:, :, i] + du_i = (d_kv_i * k_i[..., None] * v_i[..., None, :]).sum(-1) + du += du_i.sum(0) + dk_i = (dh * v_i[..., None, :]).sum(-1) + dk_aux[:, :, i] = dk_i + dk_i += (d_kv_i * u[None, ..., None] * v_i[..., None, :]).sum(-1) + dv_i = (d_kv_i * u[None, ..., None] * k_i[..., None]).sum(-2) + dv_i += (dh * k_i[..., None]).sum(-2) + + dk[:, :, i] = dk_i + dv[:, :, i] = dv_i + dh = dh * w[:, :, i, :, None].exp() + d_kv_i + + # dw = q * dq_aux - k * dk_aux + dw = torch.zeros_like(w) + for i in range(T - 2, -1, -1): + dw[:, :, i] = dw[:, :, i+1] + dq_aux[:, :, i+1] * q[:, :, i+1] - dk_aux[:, :, i] * k[:, :, i] + + return dq, dk, dv, dw, du, dh diff --git a/fla/ops/rwkv7/RWKV7(Goose).md b/fla/ops/rwkv7/RWKV7(Goose).md new file mode 100644 index 0000000000000000000000000000000000000000..b432f81c0bc4e363e3e7974033b7d59d942e1a52 --- /dev/null +++ b/fla/ops/rwkv7/RWKV7(Goose).md @@ -0,0 +1,603 @@ +# RWKV7 (Goose) Mechanism: Mathematical Derivation + +Zhiyuan Li + +>Special thanks to [Sonta](https://github.com/sustcsonglin) and [Beortust](https://github.com/Beortext), Sonta pointed out the correct notation for the outer product in the formulas, and Beortust corrected a considerable number of typos and also helped to improve the formatting. + +## Introduction to RWKV-7 Architecture + +RWKV-7 employs **Dynamic State Evolution** that transcends the fundamental TC0 expressivity limitations of attention/linear attention paradigms. RWKV-7 possesses NC1 expressivity, allowing it to solve many problems that attention mechanisms cannot. + +In simple terms, traditional attention mechanisms (like Transformer's QKV-softmax-attention) store multiple $\{k,v\}$ (key and value vector pairs), matching queries ($q$ alias named $r$ in RWKV) against keys to retrieve corresponding values. + +RWKV-7 takes a different approach - rather than directly storing $\{k,v\}$ pairs, it dynamically updates a state by learning relationships between keys and values from context. This updated state then processes new input queries ($q$, or $r$ in RWKV terminology) to produce outputs[^1]. + +[^1]: For a more detailed explanation of this approach, see the original article by the RWKV author: https://mp.weixin.qq.com/s/kC_Z3vuQ5B4PiRwZVeIvHQ + +Specifically, RWKV-7 maintains an internal model $v \approx k^{\top} S$. It aims to fit a simple objective: for given vector sequences $\{k\}$ and $\{v\}$, use state $S$ to transform $k_i$ into $v_i$, making the output $v$ as close as possible to the target $v$. + +For clarity on dimensions: + +$S_t \in \mathbb{R}^{d_v \times d_k}$ is the state matrix + +$k_t \in \mathbb{R}^{d_k}$ is the key vector + +$v_t \in \mathbb{R}^{d_v}$ is the value vector + +$q_t \in \mathbb{R}^{d_k}$ is the query vector (named $r$ in RWKV terminology) + +To achieve this, during inference with an L2 loss function $L=\frac{1}{2} \left\Vert v − k^{\top} S \right\Vert^2$, RWKV-7 automatically simulates dynamic gradient descent to continuously train its internal model $v \approx k^{\top} S$. + +The gradient of the L2 loss function with respect to the state matrix $S$ is: $\frac{\partial L}{\partial S} = S k k^{\top} - v k^{\top}$ + +Applying stochastic gradient descent (SGD) with this gradient yields a recurrent update formula that forms the foundation of RWKV-7's mechanism. In standard SGD, we would update the parameters by subtracting the gradient scaled by a learning rate: + +$$ +S_t = S_{t-1} - \eta_t \cdot \frac{\partial L}{\partial S} , \text{ where } L=L_t \quad S=S_{t-1} +$$ + +Incorporating weight decay factors $d_t = \exp(-\exp(w_t))$ as a form of time-dependent regularization and learning rate $\eta_t$, the gradient descent update becomes: + +$$S_t = S_{t-1} \text{Diag}(d_t) - \eta_t \cdot (S_{t-1} k_t k_t^{\top} - v_t k_t^{\top})$$ + +This can be expanded and rearranged as follows: + +$$S_t = S_{t-1} \text{Diag}(d_t) - \eta_t \cdot S_{t-1} k_t k_t^{\top} + \eta_t \cdot v_t k_t^{\top}$$ + +For notational simplicity, we denote $\text{Diag}(d_t)$ as $D_t$ (the diagonal decay matrix): + +$$S_t = S_{t-1} D_t - \eta_t \cdot S_{t-1} k_t k_t^{\top} + \eta_t \cdot v_t k_t^{\top}$$ + +In the full RWKV-7 implementation, this update rule is generalized through several key transformations: + +1. The diagonal decay term $D_t$ remains as a component-wise multiplication with $S_{t-1}$ + +2. The term $-\eta_t \cdot k_t k_t^{\top}$ is generalized to $\alpha_t \beta_t^{\top}$, where: + + - $\alpha_t$ can be initialized as $-k_t$ + - $\beta_t$ can be initialized as $\eta_t \cdot k_t$ + +3. The term $-\eta_t \cdot S_{t-1} k_t k_t^{\top}$ can be factorized and computed efficiently: + + - First compute $u_t = S_{t-1} k_t$ (matrix-vector product) + - Then compute $-\eta_t \cdot u_t k_t^{\top}$ (scaled outer product) + +4. The term $\eta_t \cdot v_t k_t^{\top}$ is directly implemented as the outer product between the value vector $v_t$ and key vector $k_t$, resulting in a rank-1 update matrix + +This leads to the final recurrence equation[^2]: + +$$ +S_t = S_{t-1} D_t + S_{t-1} \alpha_t \beta_t^{\top} + v_t k_t^{\top} \in \mathbb{R}^{d_v \times d_k} +$$ + +The output at each timestep is computed as: +$o_t = S_t r_t$ + +Where $r_t \in \mathbb{R}^{d_k}$ is the query vector (named $r$ in RWKV terminology), typically scaled by a factor of $\frac{1}{\sqrt{d_k}}$. This formulation allows RWKV-7 to continuously adapt its internal representation based on context, transcending the limitations of traditional attention mechanisms. + +[^2]: For a more detailed explanation, see the triton codes. Note: In the optimized Triton implementation, `w` is already the log of the decay factor, so there's only one exponential operation needed. https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/rwkv7/fused_recurrent.py#L94 + +This formulation allows more flexibility in how the state evolves while maintaining the core gradient descent learning dynamics. + +## 1. Forward Pass Recurrence Equation + +In the implementation, the state update is defined as: + +For each batch (bi) and head (hi), at time step t: + +```python +w_t = torch.exp(-torch.exp(w[bi, hi, t])) # shape [K] +sa = (state[bi, hi] * a_t[None, :]).sum(dim=1) # shape [V] +state[bi, hi] = w_t[None, :] * state[bi, hi] + sa[:, None] * b_t[None, :] + k_t[None, :] * v_t[:, None] +``` + +Where state[bi, hi] has shape [V, K], representing a state matrix that maps from K-dimensional keys to V-dimensional values. + +## 2. Backward Pass Derivation + +### 2.1 Gradient of Loss w.r.t. State + +For time step t, if L is the loss function, dstate_curr = ∂L/∂state[bi, hi, t+1] is the gradient of the current state: + +``` +dstate_curr = dstate[bi, hi] + q_t[None, :] * doutput[bi, hi, t][:, None] +``` + +This includes gradients propagated from future time steps dstate[bi, hi] and gradients from the current output. + +### 2.2 Gradient w.r.t. Query q_t + +``` +dq[bi, hi, t] = torch.matmul(doutput[bi, hi, t], curr_state) * scale +``` + +### 2.3 Gradient w.r.t. Decay Parameter w_t + +For the gradient of w_t, we need to consider how it affects the state update: + +1. For the `w_t[None, :] * state[bi, hi]` component of the state update: + +First, compute the derivative of L with respect to w_t: + +``` +∂L/∂w_t[k] = ∑_v (dstate_curr[v,k] * prev_state[v,k]) +``` + +This equation sums over the v dimension for each position k, resulting in a vector of shape [K]. + +Then, compute the derivative of w_t with respect to w: + +``` +∂w_t[k]/∂w[k] = -exp(w[k]) * exp(-exp(w[k])) = -exp(w[k]) * w_t[k] +``` + +Finally, apply the chain rule: + +``` +∂L/∂w[k] = ∂L/∂w_t[k] * ∂w_t[k]/∂w[k] + = (∑_v dstate_curr[v,k] * prev_state[v,k]) * (-exp(w[k]) * w_t[k]) +``` + +In code, this is expressed as: + +```python +dw[bi, hi, t] += -torch.sum(dstate_curr * prev_state, dim=0) * torch.exp(w[bi, hi, t]) * w_t +``` + +Or equivalently: + +```python +dw[bi, hi, t] += -torch.sum(dstate_curr * prev_state, dim=0) * torch.exp(w[bi, hi, t]) * torch.exp(-torch.exp(w[bi, hi, t])) +``` + +### 2.4 Gradient w.r.t. k_t and v_t + +For the `k_t[None, :] * v_t[:, None]` component: + +```python +dk[bi, hi, t] += torch.sum(dstate_curr * v_t[:, None], dim=0) +dv[bi, hi, t] += torch.sum(dstate_curr * k_t[None, :], dim=1) +``` + +### 2.5 Gradient w.r.t. α_t and β_t (a_t and b_t in code) + +For the `sa[:, None] * b_t[None, :]` component, where `sa = (state[bi, hi] * a_t[None, :]).sum(dim=1)`: + +```python +db[bi, hi, t] += torch.sum(dstate_curr * sa[:, None], dim=0) +dsa = torch.sum(dstate_curr * b_t[None, :], dim=1) +da[bi, hi, t] += torch.sum(prev_state * dsa[:, None], dim=0) +``` + +### 2.6 Gradient w.r.t. Previous State S\_{t-1} + +Finally, we compute the gradient of the previous state for backpropagation: + +```python +dstate_from_sa = a_t[None, :] * dsa[:, None] +dstate_from_decay = dstate_curr * w_t[None, :] +dstate[bi, hi] = dstate_from_sa + dstate_from_decay +``` + +```python +# -*- coding: utf-8 -*- +from typing import Optional, Tuple + +import torch + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def naive_recurrent_rwkv7( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + a: torch.Tensor, # Dynamic learning rate modulator + b: torch.Tensor, # State update modulator + scale: float = 1.0, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = True, +): + """ + Naive recurrent implementation of RWKV-7 (Goose) attention mechanism. + Modified from bo's code. + https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/rwkv_v7_demo.py#L170 + + Args: + q, k, v: Query, Key, and Value tensors + w: Time decay weights + a: Dynamic learning rate modulator, influences the in-context learning rate + b: State update modulator, directly participates in state update calculation + scale: Scaling factor for attention scores + initial_state: Initial state for the recurrent computation + output_final_state: Whether to output the final state + + Returns: + Attention output and optionally the final state + """ + torch_dtype = q.dtype if q.dtype in [torch.float64, torch.float] else torch.float + orig_dtype = q.dtype + B, H, L, N, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1] + q, k, v, w, a, b = (x.to(dtype=torch_dtype) for x in (q, k, v, w, a, b)) + # q, k, v, a, b, w, + # shape: (B, H, L, D), (B, H, L, D), (B, H, T, V), (B, H, L, D), (B, H, L, D), (B, H, L, D) + state = torch.zeros(B, H, V, N, dtype=torch_dtype, device=q.device) + o = torch.zeros_like(v) + + if scale == -1.0: + scale = N ** -0.5 + + if initial_state is not None: + state += initial_state.to(dtype=torch_dtype) + + for t in range(L): + q_t = q[:, :, t] * scale + k_t = k[:, :, t] + v_t = v[:, :, t] + a_t = a[:, :, t] + b_t = b[:, :, t] + + # from bo's code + sab = torch.einsum('bhik,bhk,bhj->bhij', state, a_t, b_t) + state = state * torch.exp(-torch.exp(w[:, :, t, None, :])) + sab + torch.einsum('bhj,bhi->bhij', k_t, v_t) + o[:, :, t] = torch.einsum('bhj,bhij->bhi', q_t, state) + + if not output_final_state: + ht = None + elif initial_state is not None: + ht = state.to(initial_state.dtype) + else: + ht = state.to(orig_dtype) + + return o.to(orig_dtype), ht + + +def naive_recurrent_rwkv7_2( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + a: torch.Tensor, # Dynamic learning rate modulator + b: torch.Tensor, # State update modulator + scale: float = 1.0, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = True, +): + """ + Naive recurrent implementation of RWKV-7 (Goose) attention mechanism. + + Args: + q, k, v: Query, Key, and Value tensors + w: Time decay weights + a: Dynamic learning rate modulator, influences the in-context learning rate + b: State update modulator, directly participates in state update calculation + scale: Scaling factor for attention scores + initial_state: Initial state for the recurrent computation + output_final_state: Whether to output the final state + + Returns: + Attention output and optionally the final state + """ + torch_dtype = q.dtype if q.dtype in [torch.float64, torch.float] else torch.float + orig_dtype = q.dtype + B, H, L, N, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1] + q, k, v, w, a, b = (x.to(dtype=torch_dtype) for x in (q, k, v, w, a, b)) + # q, k, v, a, b, w, + # shape: (B, H, L, D), (B, H, L, D), (B, H, T, V), (B, H, L, D), (B, H, L, D), (B, H, L, D) + state = torch.zeros(B, H, V, N, dtype=torch_dtype, device=q.device) + o = torch.zeros_like(v) + + if scale == -1.0: + scale = N ** -0.5 + + if initial_state is not None: + state += initial_state.to(dtype=torch_dtype) + + for t in range(L): + for bi in range(B): + for hi in range(H): + q_t = q[bi, hi, t] * scale + k_t = k[bi, hi, t] + v_t = v[bi, hi, t] + a_t = a[bi, hi, t] + b_t = b[bi, hi, t] + w_t = torch.exp(-torch.exp(w[bi, hi, t])) + + # h: [V, K], a_t [K] -> [1, K] + # sa: [V] + sa = (state[bi, hi] * a_t[None, :]).sum(dim=1) + + state[bi, hi] = w_t[None, :] * state[bi, hi] + sa[:, None] * b_t[None, :] + k_t[None, :] * v_t[:, None] + y = (state[bi, hi] * q_t[None, :]).sum(dim=1) + + o[bi, hi, t] = y + + ht = state if output_final_state else None + return o.to(orig_dtype), ht + + +@torch.no_grad() +def naive_recurrent_rwkv7_2_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + doutput: torch.Tensor, + dh_t: Optional[torch.Tensor] = None, + scale: float = 1.0, + dtype: Optional[torch.dtype] = None +): + """ + Backward pass for the naive_recurrent_rwkv7_2 implementation. + + Args: + q, k, v, w, a, b: Original forward pass inputs + doutput: Gradient of the loss with respect to the output + dh_t: Gradient of the loss with respect to the final state (if any) + scale: Scaling factor used in the forward pass + dtype: Optional dtype for computation + + Returns: + Gradients with respect to all inputs + """ + torch_dtype = q.dtype if q.dtype in [torch.float64, torch.float] else torch.float + q, k, v, w, a, b, doutput = (x.to(dtype=torch_dtype) for x in (q, k, v, w, a, b, doutput)) + if dh_t is not None: + dh_t = dh_t.to(dtype=torch_dtype) + + B, H, L, N, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1] + + # Initialize gradients + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dw = torch.empty_like(w) + da = torch.empty_like(a) + db = torch.empty_like(b) + + # Initialize state gradients + dstate = torch.zeros(B, H, V, N, dtype=torch_dtype, device=q.device) + if dh_t is not None: + dstate += dh_t + + if scale == -1.0: + scale = N ** -0.5 + + # First rebuild all states from forward pass + states = [] + state = torch.zeros(B, H, V, N, dtype=torch_dtype, device=q.device) + states.append(state.clone()) + + # In practice, we don't recompute all states from the beginning. + # Instead, we use checkpointing: we save states at regular intervals (e.g., every 16 tokens) + # during the forward pass, then reconstruct intermediate states during the backward pass + # by working backwards from the nearest checkpoint. + # + # For example, to get state[t-1] from state[t]: + # state[t-1] = (state[t] - (sa * b_t + k_t * v_t)) / w_t + # + # This approach balances memory usage and computational efficiency: + # - Reduces memory by not storing every state + # - Maintains numerical stability by limiting the number of backward steps from each checkpoint + # - Allows efficient gradient computation without recomputing the entire sequence + for t in range(L): + for bi in range(B): + for hi in range(H): + q_t = q[bi, hi, t] * scale + k_t = k[bi, hi, t] + v_t = v[bi, hi, t] + a_t = a[bi, hi, t] + b_t = b[bi, hi, t] + w_t = torch.exp(-torch.exp(w[bi, hi, t])) + + sa = (state[bi, hi] * a_t[None, :]).sum(dim=1) + + state[bi, hi] = w_t[None, :] * state[bi, hi] + sa[:, None] * b_t[None, :] + k_t[None, :] * v_t[:, None] + states.append(state.clone()) + + # Backward pass through time + for t in range(L-1, -1, -1): + for bi in range(B): + for hi in range(H): + q_t = q[bi, hi, t] * scale + k_t = k[bi, hi, t] + v_t = v[bi, hi, t] + a_t = a[bi, hi, t] + b_t = b[bi, hi, t] + w_scalar = w[bi, hi, t] + w_exp = torch.exp(w_scalar) + w_t = torch.exp(-w_exp) + + curr_state = states[t+1][bi, hi] # State after update [V, K] + prev_state = states[t][bi, hi] # State before update [V, K] + + dq[bi, hi, t] = (doutput[bi, hi, t][:, None] * curr_state).sum(dim=0) * scale + + dstate_from_out = q_t[None, :] * doutput[bi, hi, t][:, None] # [V, K] + + dstate_curr = dstate[bi, hi] + dstate_from_out + + sa = (prev_state * a_t[None, :]).sum(dim=1) # [V] + + # state[bi, hi] = w_t[None, :] * prev_state + ... + dw[bi, hi, t] = -torch.sum(dstate_curr * prev_state, dim=0) * \ + w_t * w_exp + + # k_t[None, :] * v_t[:, None] -> [V, K] + dk[bi, hi, t] = torch.sum(dstate_curr * v_t[:, None], dim=0) + dv[bi, hi, t] = torch.sum(dstate_curr * k_t[None, :], dim=1) + + # sa[:, None] * b_t[None, :] -> [V, K] + db[bi, hi, t] = torch.sum(dstate_curr * sa[:, None], dim=0) + dsa = torch.sum(dstate_curr * b_t[None, :], dim=1) # [V] + + # sa = (prev_state * a_t[None, :]).sum(dim=1) + da[bi, hi, t] = torch.sum(prev_state * dsa[:, None], dim=0) + dstate_from_sa = a_t[None, :] * dsa[:, None] # [V, K] + + # w_t[None, :] * prev_state + dstate_from_decay = dstate_curr * w_t[None, :] # [V, K] + + dstate[bi, hi] = dstate_from_sa + dstate_from_decay + + return dq, dk, dv, dw, da, db, dstate + + +class NativeRecurrentRWKV7Function(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, w, a, b, scale, initial_state, + training: bool = True, dtype: Optional[torch.dtype] = None, + state_ckpt_interval: int = 16): + o, ht = naive_recurrent_rwkv7_2(q, k, v, w, a, b, scale=scale, initial_state=initial_state) + if training: + ctx.save_for_backward(q, k, v, w, a, b) + ctx.scale = scale + ctx.dtype = dtype + ctx.ckpt_interval = state_ckpt_interval + ctx.use_initial_state = initial_state is not None + return o, ht + + @staticmethod + @autocast_custom_bwd + def backward(ctx, do, dht): + q, k, v, w, a, b = ctx.saved_tensors + dq, dk, dv, dw, da, db, dh = naive_recurrent_rwkv7_2_bwd( + q, k, v, w, a, b, do, dht, ctx.scale, dtype=ctx.dtype) + dh = dh if ctx.use_initial_state else None + return dq, dk, dv, dw, da, db, None, dh, None, None + + +def recurrent_rwkv7( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = 1.0, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + r (torch.Tensor): + r of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + k (torch.Tensor): + k of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + v (torch.Tensor): + v of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + w (torch.Tensor): + decay of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`, kernel + will apply log_w = -torch.exp(w) + log_w (torch.Tensor): + log decay of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + scale (float): + scale of the attention. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (bool): + whether to use head first. Recommended to be False to avoid extra transposes. + """ + assert cu_seqlens is None + assert head_first is True + assert w is not None + if scale == -1.0: + scale = q.shape[-1] ** -0.5 + o, final_state = NativeRecurrentRWKV7Function.apply(q, k, v, w, a, b, scale, initial_state) + + return o, final_state + + +def test_autograd_function(): + """Test the custom autograd function implementation""" + # Set random seed for reproducibility + torch.manual_seed(42) + + # Define test dimensions + B, H, T, D = 1, 1, 128, 64 + V = N = D + device = 'cpu' + dtype = torch.float64 + + # Create random test inputs + q = torch.empty(B, H, T, D, device=device).uniform_(-8, 8).to(dtype=dtype).requires_grad_(True) + k = torch.empty(B, H, T, D, device=device).uniform_(-8, 8).to(dtype=dtype).requires_grad_(True) + v = torch.empty(B, H, T, D, device=device).uniform_(-8, 8).to(dtype=dtype).requires_grad_(True) + w = torch.empty(B, H, T, D, device=device).uniform_(-8, -6).to(dtype=dtype).requires_grad_(True) + + kk = torch.empty(B, H, T, D, device=device).uniform_(-8, 8) + kk = torch.nn.functional.normalize(kk, dim=-1).to(dtype=dtype) + + a = -kk.clone().requires_grad_(True) # -kk + a_scale = torch.empty(B, H, T, D, device=device).uniform_(0, 0.1).to(dtype=dtype) + b = (kk * a_scale).requires_grad_(True) # kk*a + + # Create initial state + initial_state = torch.zeros(B, H, V, N).to(torch.float64) + + # Clone inputs for the two paths we're testing + q1, k1, v1, w1, a1, b1 = q.clone().detach().requires_grad_(True), k.clone().detach().requires_grad_(True), v.clone().detach().requires_grad_( + True), w.clone().detach().requires_grad_(True), a.clone().detach().requires_grad_(True), b.clone().detach().requires_grad_(True) + q2, k2, v2, w2, a2, b2 = q.clone().detach().requires_grad_(True), k.clone().detach().requires_grad_(True), v.clone().detach().requires_grad_( + True), w.clone().detach().requires_grad_(True), a.clone().detach().requires_grad_(True), b.clone().detach().requires_grad_(True) + + # Path 1: Using naive implementation with autograd + + output1, state1 = naive_recurrent_rwkv7(q1, k1, v1, w1, a1, b1, initial_state=initial_state.clone()) + + output2, state2 = recurrent_rwkv7(q2, k2, v2, w2, a2, b2, 1.0, initial_state.clone()) + + # Check forward pass equivalence + output_diff = torch.max(torch.abs(output1 - output2)).item() + state_diff = torch.max(torch.abs(state1 - state2)).item() + + print(f"\nAutograd Function test (forward):") + print(f" Max output difference: {output_diff:.6e}") + print(f" Max state difference: {state_diff:.6e}") + + # Create loss function to test backward pass + def compute_loss(output, state): + return output.sum() # + state.sum() + + # Compute loss and gradients for both paths + loss1 = compute_loss(output1, state1) + loss1.backward() + + loss2 = compute_loss(output2, state2) + loss2.backward() + + # Compare gradients + grad_diffs = { + 'q': torch.max(torch.abs(q1.grad - q2.grad)).item(), + 'k': torch.max(torch.abs(k1.grad - k2.grad)).item(), + 'v': torch.max(torch.abs(v1.grad - v2.grad)).item(), + 'w': torch.max(torch.abs(w1.grad - w2.grad)).item(), + 'a': torch.max(torch.abs(a1.grad - a2.grad)).item(), + 'b': torch.max(torch.abs(b1.grad - b2.grad)).item(), + } + + print(f"\nAutograd Function test (backward):") + for param, diff in grad_diffs.items(): + print(f" Max {param} gradient difference: {diff:.6e}") + + +test_autograd_function() +``` diff --git a/fla/ops/rwkv7/__init__.py b/fla/ops/rwkv7/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05c58664f765449d61238ce566714a82a219811a --- /dev/null +++ b/fla/ops/rwkv7/__init__.py @@ -0,0 +1,9 @@ + +from .chunk import chunk_rwkv7 +from .fused_recurrent import fused_mul_recurrent_rwkv7, fused_recurrent_rwkv7 + +__all__ = [ + 'chunk_rwkv7', + 'fused_recurrent_rwkv7', + 'fused_mul_recurrent_rwkv7', +] diff --git a/fla/ops/rwkv7/channel_mixing.py b/fla/ops/rwkv7/channel_mixing.py new file mode 100644 index 0000000000000000000000000000000000000000..5f30b67410d00d167bd96b851b40986f29830885 --- /dev/null +++ b/fla/ops/rwkv7/channel_mixing.py @@ -0,0 +1,334 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import logging + +import torch +import triton +import triton.language as tl + +from fla.utils import ( + USE_CUDA_GRAPH, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_pytorch_version, + input_guard, +) + +logger = logging.getLogger(__name__) + +if not check_pytorch_version('2.4'): + logger.warning('PyTorch < 2.4 detected - computations may be slower due to lack of optimizations') + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': block_size}) + for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] + ], + key=['hidden_dim'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def rwkv_seq_mix_kernel( + x_ptr, + x_prev_ptr, + mix_k_ptr, + output_ptr, + batch_size: tl.constexpr, + token_length, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + block_start = tl.program_id(0) * BLOCK_SIZE + block_idx = block_start + tl.arange(0, BLOCK_SIZE)[:] + + total_seq_dim = token_length * hidden_dim + batch_idx = block_idx // total_seq_dim + seq_and_feat = block_idx % total_seq_dim + seq_idx = seq_and_feat // hidden_dim + feat_idx = seq_and_feat % hidden_dim + + is_valid = (batch_idx < batch_size) & (seq_idx < token_length) + + x_idx = batch_idx * total_seq_dim + seq_idx * hidden_dim + feat_idx + + curr_x = tl.load(x_ptr + x_idx, mask=is_valid, other=0.0).to(tl.float32) + k_value = tl.load(mix_k_ptr + feat_idx).to(tl.float32) + + is_first = seq_idx < 1 + prev_state_idx = batch_idx * hidden_dim + feat_idx + prev_state = tl.load(x_prev_ptr + prev_state_idx, + mask=(is_first & is_valid), + other=0.0).to(tl.float32) + + prev_x_idx = x_idx - hidden_dim + prev_x = tl.load(x_ptr + prev_x_idx, + mask=(~is_first & is_valid), + other=0.0).to(tl.float32) + + prev_value = tl.where(is_first, prev_state, prev_x) + state_diff = prev_value - curr_x + mixed = state_diff * k_value + result = tl.cast(curr_x + mixed, dtype=output_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(output_ptr + x_idx, result, mask=is_valid) + + +@triton.jit +def rwkv_channel_mixing_pow_and_relu( + in_ptr, + out_ptr, + BLOCK_SIZE: tl.constexpr, +): + """Fused ReLU and Power operation: x = ReLU(x)^2""" + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE) + x0 = xindex + x = tl.load(in_ptr + (x0), None) + x = tl.maximum(x, 0.0).to(tl.float32) + x = tl.cast(x * x, dtype=out_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(out_ptr + (x0), x, None) + + +def rwkv_mix_torch(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor): + if x_prev.dim() == 2: + x_prev = x_prev.unsqueeze(1) # (batch_size, 1, hidden_dim) + xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x + k = x.addcmul(xx, x_k) + return k + + +def rwkv_relu_and_square_torch(x: torch.Tensor): + return torch.relu(x) ** 2 + + +def rwkv_mix_fwd(x, x_prev, x_k): + has_batch = x.dim() == 3 + + if has_batch: + batch_size, token_length, hidden_dim = x.shape + else: + token_length, hidden_dim = x.shape + batch_size = 1 + x = x.unsqueeze(0) + x_prev = x_prev.unsqueeze(0) + + token_length = x.shape[1] + hidden_dim = x.shape[2] + total_elements = batch_size * token_length * hidden_dim + + output = torch.empty_like(x) + + def grid(meta): return ( + (total_elements + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 + 1, # grid_1 + 1, # grid_2 + ) + + rwkv_seq_mix_kernel[grid]( + x.contiguous(), + x_prev.contiguous(), + x_k.squeeze(), + output, + batch_size=batch_size, + token_length=token_length, + hidden_dim=hidden_dim, + ) + if not has_batch: + output = output.squeeze(0) + return output + + +def rwkv_relu_and_square_fwd(x: torch.Tensor, inplace: bool = True): + """ + Triton implementation of RWKV's ReLU and square operation + Args: + x: Input tensor + Returns: + Tensor after ReLU and square operations + """ + x = x.contiguous() + output = x if inplace else torch.empty_like(x) + + def grid(meta): return ( + (output.numel() + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 + 1, # grid_1 + 1, # grid_2 + ) + rwkv_channel_mixing_pow_and_relu[grid]( + x, + output, + BLOCK_SIZE=4096, + ) + + return output + + +@triton.jit +def relu_square_bwd_kernel( + out_ptr, + forward_input_ptr, + BLOCK_SIZE: tl.constexpr, +): + """ReLU(x)^2 backward kernel + grad_input = grad_output * 2 * x if x > 0 else 0 + """ + pid = tl.program_id(0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + x = tl.load(forward_input_ptr + offsets).to(tl.float32) + grad = tl.load(out_ptr + offsets).to(tl.float32) + + x = tl.maximum(x, 0.0) + + grad_input = grad * 2 * x + + tl.store(out_ptr + offsets, grad_input.to(out_ptr.dtype.element_ty)) + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': block_size}) + for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] + ], + key=['hidden_dim'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def rwkv_mix_bwd_kenel( + dk1_ptr0, + xk_ptr, + dx_ptr, + dx_prev_ptr, + batch_size, + token_length, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + + batch_idx = offsets // (token_length * hidden_dim) + seq_feat = offsets % (token_length * hidden_dim) + seq_idx = seq_feat // hidden_dim + feat_idx = seq_feat % hidden_dim + + is_valid = offsets < (batch_size * token_length * hidden_dim) + + dk1 = tl.load(dk1_ptr0 + offsets, mask=is_valid) + xk = tl.load(xk_ptr + feat_idx, mask=is_valid) + prod = dk1 * xk + + mask_next = seq_idx < (token_length - 1) + next_offset = offsets + hidden_dim + dk1_next = tl.load(dk1_ptr0 + next_offset, mask=mask_next & is_valid, other=0.0) + prod_next = dk1_next * xk + dx_val = dk1 - prod + tl.where(mask_next, prod_next, 0.0) + dx_val = tl.cast(dx_val, dtype=dx_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(dx_ptr + offsets, dx_val, mask=is_valid) + + dx_prev_offset = batch_idx * hidden_dim + feat_idx + is_first_step = seq_idx == 0 + + tl.store( + dx_prev_ptr + dx_prev_offset, + tl.cast(prod, dtype=dx_prev_ptr.dtype.element_ty), + mask=is_first_step, + ) + + +@torch.compile(fullgraph=True) +def compute_x_k_grad(dk1, x, x_prev): + """ + Args: + dk1: (batch*seq_len, hidden_dim) + x: (batch, seq_len, hidden_dim) + x_prev: (batch, hidden_dim) or (batch, 1, hidden_dim) + """ + + if x_prev.dim() == 2: + x_prev = x_prev.unsqueeze(1) # (batch, 1, hidden_dim) + xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x # (batch, seq_len, hidden_dim) + + # (hidden_dim,) --> (1, 1, hidden_dim) + grad_x_k = (dk1 * xx.reshape(-1, x.shape[2])).sum(dim=0).view(1, 1, -1) + return grad_x_k + + +def rwkv_channel_mixing_bwd(grad_output, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, inplace=True): + batch_size = x.shape[0] if x.dim() == 3 else 1 + seq_len, n_embd = x.shape[-2], x.shape[-1] + + dV = k.transpose(-2, -1) @ grad_output + dk = grad_output @ value_weight.transpose(-2, -1) + + BLOCK_SIZE = 4096 + grid = ((dk.numel() + BLOCK_SIZE - 1) // BLOCK_SIZE,) + relu_square_bwd_kernel[grid]( + dk, + k1_K, + BLOCK_SIZE=BLOCK_SIZE, + ) + + dK = k1.transpose(-2, -1) @ dk + dk1 = dk @ key_weight.transpose(-2, -1) + dk1 = dk1.view(-1, n_embd).contiguous() + + dk_reduced = compute_x_k_grad(dk1, x, x_prev) + dx_prev = torch.empty_like(x_prev) if not inplace else x_prev + dx = torch.empty_like(x) if not inplace else x + + def grid(meta): return ((batch_size * seq_len * n_embd + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], 1, 1) + rwkv_mix_bwd_kenel[grid]( + dk1, + x_k.squeeze(), + dx, + dx_prev, + batch_size, + seq_len, + n_embd, + ) + # dx_prev.shape batch_size, seq_len, n_embd + return dx, dx_prev, dk_reduced, dK, dV + + +class Rwkv7ChannelMixing(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, x, x_prev, x_k, key_weight, value_weight, inplace: bool = True): + k1 = rwkv_mix_fwd(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_fwd(k1_K, inplace=True) + ctx.save_for_backward(x, x_prev, x_k, key_weight, value_weight) + ctx.inplace = inplace + return k @ value_weight + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dkv): + x, x_prev, x_k, key_weight, value_weight = ctx.saved_tensors + k1 = rwkv_mix_fwd(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_fwd(k1_K, inplace=False) + dx, dx_prev, dk_reduced, dK, dV = rwkv_channel_mixing_bwd( + dkv, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, ctx.inplace) + return dx, dx_prev, dk_reduced, dK, dV, None + + +def channel_mixing_rwkv7(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor, + key_weight: torch.Tensor, value_weight: torch.Tensor, inplace: bool = True): + assert x.dim() == 3 + + return Rwkv7ChannelMixing.apply(x, x_prev, x_k, key_weight, value_weight, inplace), x[-1, :] + + +def channel_mixing_rwkv7_torch(x, x_prev, x_k, key_weight, value_weight): + k1 = rwkv_mix_torch(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_torch(k1_K) + return k @ value_weight, x[-1, :] diff --git a/fla/ops/rwkv7/chunk.py b/fla/ops/rwkv7/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..8f48072e14e6b3dd116924936914df218b02e76c --- /dev/null +++ b/fla/ops/rwkv7/chunk.py @@ -0,0 +1,91 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch + +from fla.ops.generalized_delta_rule import chunk_dplr_delta_rule + + +def chunk_rwkv7( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = 1.0, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, + safe_gate: bool = False, + chunk_size: int | None = None, +): + """ + Args: + r (torch.Tensor): + r of shape `[B, T, H, K]`. + w (torch.Tensor): + log decay of shape `[B, T, H, K]`. + k (torch.Tensor): + k of shape `[B, T, H, K]`. + v (torch.Tensor): + v of shape `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, T, H, K]`. + scale (float): + scale of the attention. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cu_seqlens_cpu (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + safe_gate (bool): + Whether the kernel can assume the input gate values `g` are in a safe range. + When `True`, the kernel can use M=16 TensorCore acceleration. + The safe range is approximately [-5, 0). Default: `False`. + chunk_size (Optional[int]): + Chunk size for the chunked computation. Default: `None`, which means 16. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and r.shape[1] < r.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + return chunk_dplr_delta_rule( + q=r, + k=k, + v=v, + a=a, + b=b, + gk=w, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + safe_gate=safe_gate, + chunk_size=chunk_size, + head_first=head_first, + ) diff --git a/fla/ops/rwkv7/fused_addcmul.py b/fla/ops/rwkv7/fused_addcmul.py new file mode 100644 index 0000000000000000000000000000000000000000..048df6d43acbd3ecbace1af05c7698f0a14be03d --- /dev/null +++ b/fla/ops/rwkv7/fused_addcmul.py @@ -0,0 +1,295 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import logging +import os +import sys + +import torch +import triton +import triton.language as tl +from packaging.version import Version + +from fla.utils import IS_AMD, USE_CUDA_GRAPH, autotune_cache_kwargs, check_pytorch_version, input_guard + +logger = logging.getLogger(__name__) + +if not check_pytorch_version('2.4'): + logger.warning('PyTorch < 2.4 detected - computations may be slower due to lack of optimizations') + + +def identity_decorator(fn): + return fn + + +current_python_version = Version(f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}") +min_torch_compile_version = Version("3.11") +fla_use_compile = os.getenv('FLA_USE_COMPILE', '1').lower() in ('1', 'true', 'yes') + +if current_python_version >= min_torch_compile_version and fla_use_compile: + torch_compile = torch.compile(fullgraph=True) +else: + logger.warning('torch.compile is not available in Python 3.10, using identity decorator instead') + torch_compile = identity_decorator + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + for BT in [2, 4, 8] + ], + key=['BD'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def fused_addcmul_fwd_kernel( + hidden, + delta, + ixr, ixw, ixk, ixv, ixa, ixg, + oxr, oxw, oxk, oxv, oxa, oxg, + use_xg: tl.constexpr, + T, + T_OFFSET, + BT: tl.constexpr, + D: tl.constexpr, + BD: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) * BT + + bos = i_b * (T + T_OFFSET) + t_vec = i_t + T_OFFSET + tl.arange(0, BT) + mask_t = t_vec < (T + T_OFFSET) + o_d = tl.arange(0, BD)[None, :] + off_vec = (bos + t_vec)[:, None] * D + o_d + m_d = o_d < D + mask = mask_t[:, None] & m_d + + b_h = tl.load(hidden + off_vec, mask=mask, other=0.) + b_x = tl.load(delta + off_vec, mask=mask, other=0.) + b_r = tl.load(ixr + o_d, mask=m_d) + b_w = tl.load(ixw + o_d, mask=m_d) + b_k = tl.load(ixk + o_d, mask=m_d) + b_v = tl.load(ixv + o_d, mask=m_d) + b_a = tl.load(ixa + o_d, mask=m_d) + + o_r = tl.fma(b_x, b_r, b_h) + o_w = tl.fma(b_x, b_w, b_h) + o_k = tl.fma(b_x, b_k, b_h) + o_v = tl.fma(b_x, b_v, b_h) + o_a = tl.fma(b_x, b_a, b_h) + + tl.store(oxr + off_vec, o_r.to(oxr.dtype.element_ty), mask=mask) + tl.store(oxw + off_vec, o_w.to(oxw.dtype.element_ty), mask=mask) + tl.store(oxk + off_vec, o_k.to(oxk.dtype.element_ty), mask=mask) + tl.store(oxv + off_vec, o_v.to(oxv.dtype.element_ty), mask=mask) + tl.store(oxa + off_vec, o_a.to(oxa.dtype.element_ty), mask=mask) + + if use_xg: + b_g = tl.load(ixg + o_d, mask=m_d) + o_g = tl.fma(b_x, b_g, b_h) + tl.store(oxg + off_vec, o_g.to(oxg.dtype.element_ty), mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS_AUTOTUNE + for num_stages in [1, 2, 3] + for BT in [2, 4, 8] + ], + key=['BD'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit +def addcmul_bwd_kernel1( + ixr, + ixw, + ixk, + ixv, + ixa, + ixg, + dxr, + dxw, + dxk, + dxv, + dxa, + dxg, + ghidden, + gx, + use_xg: tl.constexpr, + T, + T_OFFSET, + BT: tl.constexpr, + D: tl.constexpr, + BD: tl.constexpr, + DTYPE: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) + + t_idx = T_OFFSET + i_t * BT + tl.arange(0, BT)[:, None] + mask_t = t_idx < (T + T_OFFSET) + + d_idx = tl.arange(0, BD)[None, :] + mask_d = d_idx < D + mask = mask_t & mask_d + + offset_base = i_b * (T + T_OFFSET) * D + x_idx = (offset_base + t_idx * D + d_idx).to(tl.uint32) + + b_dxr = tl.load(dxr + x_idx, mask=mask).to(DTYPE) + b_dxw = tl.load(dxw + x_idx, mask=mask).to(DTYPE) + b_dxk = tl.load(dxk + x_idx, mask=mask).to(DTYPE) + b_dxv = tl.load(dxv + x_idx, mask=mask).to(DTYPE) + b_dxa = tl.load(dxa + x_idx, mask=mask).to(DTYPE) + + b_ixr = tl.load(ixr + d_idx, mask=mask_d).to(DTYPE) + b_ixw = tl.load(ixw + d_idx, mask=mask_d).to(DTYPE) + b_ixk = tl.load(ixk + d_idx, mask=mask_d).to(DTYPE) + b_ixv = tl.load(ixv + d_idx, mask=mask_d).to(DTYPE) + b_ixa = tl.load(ixa + d_idx, mask=mask_d).to(DTYPE) + + g_hidden = b_dxr + b_dxw + b_dxk + b_dxv + b_dxa + g_x = b_dxr * b_ixr + b_dxw * b_ixw + b_dxk * b_ixk + b_dxv * b_ixv + b_dxa * b_ixa + + if use_xg: + b_dxg = tl.load(dxg + x_idx, mask=mask).to(DTYPE) + b_ixg = tl.load(ixg + d_idx, mask=mask_d).to(DTYPE) + g_hidden += b_dxg + g_x += b_dxg * b_ixg + + tl.store(ghidden + x_idx, g_hidden.to(ghidden.dtype.element_ty), mask=mask) + tl.store(gx + x_idx, g_x.to(gx.dtype.element_ty), mask=mask) + + +def addcmul_bwd1(d_xr, d_xw, d_xk, d_xv, d_xa, d_xg, + x_r, x_w, x_k, x_v, x_a, x_g, hidden_states, delta, use_xg, inplace=True): + B, T, D = hidden_states.size() + g_hiddn = hidden_states if inplace else torch.empty_like(hidden_states) + g_delta = torch.empty_like(delta) + for t in range(0, T, 65536): + T_OFFSET = t + T_SIZE = min(65536, T - t) + def grid(meta): return (B, triton.cdiv(T_SIZE, meta['BT'])) + addcmul_bwd_kernel1[grid]( + ixr=x_r, + ixw=x_w, + ixk=x_k, + ixv=x_v, + ixa=x_a, + ixg=x_g, + dxr=d_xr, + dxw=d_xw, + dxk=d_xk, + dxv=d_xv, + dxa=d_xa, + dxg=d_xg, + ghidden=g_hiddn, + gx=g_delta, + use_xg=use_xg, + T=T_SIZE, + T_OFFSET=T_OFFSET, + D=D, + BD=triton.next_power_of_2(D), + DTYPE=tl.float16 if hidden_states.dtype == torch.float16 else tl.float32, + ) + return g_hiddn, g_delta + + +@torch_compile +def addcmul_bwd2(d_oxr, d_xw, d_xk, d_xv, d_xa, d_xg, delta, use_xg: bool): + g_xr = (d_oxr * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) + g_xw = (d_xw * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) + g_xk = (d_xk * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) + g_xv = (d_xv * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) + g_xa = (d_xa * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) + g_xg = (d_xg * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) if use_xg else None + return g_xr, g_xw, g_xk, g_xv, g_xa, g_xg + + +class Rwkv7FusedAddcmul(torch.autograd.Function): + @staticmethod + @input_guard + def forward( + ctx, hidden_states, delta, + x_r, x_w, x_k, x_v, x_a, x_g, + ): + B, T, D = hidden_states.size() + oxr = torch.empty_like(hidden_states) + oxw = torch.empty_like(hidden_states) + oxk = torch.empty_like(hidden_states) + oxv = torch.empty_like(hidden_states) + oxa = torch.empty_like(hidden_states) + if x_g is not None: + use_xg = True + oxg = torch.empty_like(hidden_states) + else: + use_xg = False + oxg = None + + for t in range(0, T, 65536): + T_OFFSET = t + T_SIZE = min(65536, T - t) + def grid(meta): return (B, triton.cdiv(T_SIZE, meta['BT'])) + fused_addcmul_fwd_kernel[grid]( + hidden_states, delta, + x_r, x_w, x_k, x_v, x_a, x_g, + oxr, oxw, oxk, oxv, oxa, oxg, + use_xg, + T=T_SIZE, + T_OFFSET=T_OFFSET, + D=D, + BD=triton.next_power_of_2(D), + ) + + ctx.save_for_backward(hidden_states, delta, + x_r, x_w, x_k, x_v, x_a, x_g) + ctx.use_xg = use_xg + return oxr, oxw, oxk, oxv, oxa, oxg + + @staticmethod + @input_guard + def backward(ctx, dxr, + dxw, dxk, dxv, dxa, dxg): + hidden_states, delta, x_r, x_w, x_k, x_v, x_a, x_g = ctx.saved_tensors + + d_hiddn, d_xx = addcmul_bwd1(dxr, dxw, dxk, dxv, dxa, dxg, x_r, x_w, x_k, x_v, x_a, x_g, + hidden_states, delta, ctx.use_xg) + + d_ixr, d_ixw, d_ixk, d_ixv, d_ixa, d_ixg = addcmul_bwd2(dxr, dxw, dxk, dxv, dxa, dxg, delta, ctx.use_xg) + + return d_hiddn, d_xx, d_ixr, d_ixw, d_ixk, d_ixv, d_ixa, d_ixg + + +def fused_addcmul_rwkv7( + hidden_states: torch.Tensor, + delta: torch.Tensor, + xr: torch.Tensor, + xw: torch.Tensor, + xk: torch.Tensor, + xv: torch.Tensor, + xa: torch.Tensor, + xg: torch.Tensor | None = None, +): + if hidden_states.shape[1] == 1: + # Special case for decode + return torch_addcmul_rwkv7(hidden_states, delta, xr, xw, xk, xv, xa, xg) + return Rwkv7FusedAddcmul.apply(hidden_states, delta, xr, xw, xk, xv, xa, xg) + + +def torch_addcmul_rwkv7(hidden_states, delta, xr, xw, xk, xv, xa, xg=None): + oxr = torch.addcmul(hidden_states, delta, xr) + oxw = torch.addcmul(hidden_states, delta, xw) + oxk = torch.addcmul(hidden_states, delta, xk) + oxv = torch.addcmul(hidden_states, delta, xv) + oxa = torch.addcmul(hidden_states, delta, xa) + if xg is not None: + oxg = torch.addcmul(hidden_states, delta, xg) + return oxr, oxw, oxk, oxv, oxa, oxg + else: + return oxr, oxw, oxk, oxv, oxa, None + return oxr, oxw, oxk, oxv, oxa, None diff --git a/fla/ops/rwkv7/fused_k_update.py b/fla/ops/rwkv7/fused_k_update.py new file mode 100644 index 0000000000000000000000000000000000000000..69d3c505a7e87451d9d0649ccb98f4b21209cc96 --- /dev/null +++ b/fla/ops/rwkv7/fused_k_update.py @@ -0,0 +1,351 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard + +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] + + +def k_update_ref(k: torch.Tensor, a: torch.Tensor, ka: torch.Tensor) -> torch.Tensor: + return k.addcmul(k * (a - 1), ka) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=w, num_stages=s) + for w in NUM_WARPS_AUTOTUNE + for s in [1, 2, 3] + ], + key=['BD'], + **autotune_cache_kwargs, +) +@triton.jit +def k_update_fwd_kernel_short( + k, a, ka, out, + cu_seqlens, + T, D, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_b).to(tl.int32) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + g_t = bos + i_t + if g_t >= eos: + return + offset = g_t * D + else: + g_t = i_t + offset = i_b * T * D + g_t * D + + o_d = tl.arange(0, BD) + m_d = o_d < D + off = offset + o_d + + b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) + b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) + b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) + + out_val = b_k * (1 + (b_a - 1) * b_ka) + tl.store(out + off, out_val.to(out.dtype.element_ty), mask=m_d) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=w, num_stages=s) + for w in NUM_WARPS_AUTOTUNE + for s in [1, 2, 3] + ], + key=['BD', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit +def k_update_fwd_kernel_long( + k, a, ka, out, + cu_seqlens, chunk_indices, + T, D, + BD: tl.constexpr, BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ + tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, eos - bos) + else: + bos = i_b * T + eos = (i_b + 1) * T + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, T) + + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + + for t in range(t_start, t_end): + global_t = bos + t + off = global_t * D + o_d + b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) + b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) + b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) + out_val = b_k * (1 + (b_a - 1) * b_ka) + tl.store(out + off, out_val.to(out.dtype.element_ty), mask=m_d) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=w, num_stages=s) + for w in NUM_WARPS_AUTOTUNE + for s in [1, 2, 3] + for BT in [2, 4, 8] + ], + key=['BD'], + **autotune_cache_kwargs, +) +@triton.jit +def k_update_bwd_kernel_short( + grad_out, k, a, ka, + dk, da, dka, + cu_seqlens, + T, D, + BT: tl.constexpr, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_b, i_t_base = tl.program_id(0), tl.program_id(1) * BT + + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_b).to(tl.int32) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + seq_len = eos - bos + else: + bos = i_b * T + eos = (i_b + 1) * T + seq_len = T + + t_vec = i_t_base + tl.arange(0, BT) + mask_t = t_vec < seq_len + global_t_vec = bos + t_vec + + o_d = tl.arange(0, BD)[None, :] + m_d = o_d < D + off = global_t_vec[:, None] * D + o_d + + b_go = tl.load(grad_out + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) + b_k = tl.load(k + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) + b_a = tl.load(a + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) + b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) # [1, BD] + + dk_vec = b_go * (1 + (b_a - 1) * b_ka) + da_vec = b_go * b_k * b_ka + dka_vec = b_go * b_k * (b_a - 1) + tl.store(dk + off, dk_vec.to(dk.dtype.element_ty), mask=mask_t[:, None] & m_d) + tl.store(da + off, da_vec.to(da.dtype.element_ty), mask=mask_t[:, None] & m_d) + tl.store(dka + off, dka_vec.to(dka.dtype.element_ty), mask=mask_t[:, None] & m_d) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=w, num_stages=s) + for w in NUM_WARPS_AUTOTUNE + for s in [1, 2, 3] + ], + key=['BD', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit +def k_update_bwd_kernel_long( + grad_out, k, a, ka, + dk, da, dka, + cu_seqlens, chunk_indices, + T, D, + BD: tl.constexpr, BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if IS_VARLEN: + i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ + tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, eos - bos) + else: + bos = i_b * T + eos = (i_b + 1) * T + t_start = i_t_blk * BT + t_end = tl.minimum(t_start + BT, T) + + o_d = i_d * BD + tl.arange(0, BD) + m_d = o_d < D + + for t in range(t_start, t_end): + global_t = bos + t + off = global_t * D + o_d + + b_go = tl.load(grad_out + off, mask=m_d, other=0.).to(tl.float32) + b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) + b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) + b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) + + tl.store(dk + off, (b_go * (1 + (b_a - 1) * b_ka)).to(dk.dtype.element_ty), mask=m_d) + tl.store(da + off, (b_go * b_k * b_ka).to(da.dtype.element_ty), mask=m_d) + tl.store(dka + off, (b_go * b_k * (b_a - 1)).to(dka.dtype.element_ty), mask=m_d) + + +def k_update_fwd( + k: torch.Tensor, + a: torch.Tensor, + ka: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, D = k.shape + out = torch.empty_like(k) + use_short = T <= 512 + + if use_short: + if cu_seqlens is not None: + N = len(cu_seqlens) - 1 + else: + N = B + BD = triton.next_power_of_2(D) + grid = (N, T) + k_update_fwd_kernel_short[grid]( + k, a, ka, out, + cu_seqlens, + T, D, + BD=BD, + ) + else: + BT = min(64, triton.next_power_of_2( + triton.cdiv(max(16, B * T), get_multiprocessor_count(k.device.index)), + )) + if cu_seqlens is not None: + chunk_idx = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) + NT = len(chunk_idx) + N = len(cu_seqlens) - 1 + else: + chunk_idx = None + NT = triton.cdiv(T, BT) + N = B + + BD = triton.next_power_of_2(D) + + def grid(meta): + return (triton.cdiv(D, meta['BD']), NT, N) + + k_update_fwd_kernel_long[grid]( + k, a, ka, out, + cu_seqlens, chunk_idx, + T, D, + BD=BD, BT=BT, + ) + + return out, use_short, N, T + + +def k_update_bwd( + grad_out: torch.Tensor, + k: torch.Tensor, + a: torch.Tensor, + ka: torch.Tensor, + cu_seqlens: torch.Tensor | None, + use_short: bool, + N: int, + T: int, + cu_seqlens_cpu: torch.LongTensor | None = None, +): + B, _, D = grad_out.shape + dk = torch.empty_like(k) + da = torch.empty_like(a) + dka_tmp = torch.empty_like(k, dtype=torch.float32) + + if use_short: + BD = triton.next_power_of_2(D) + def grid(meta): return (N, triton.cdiv(T, meta['BT'])) + k_update_bwd_kernel_short[grid]( + grad_out, k, a, ka, + dk, da, dka_tmp, + cu_seqlens, + T, D, + BD=BD, + ) + else: + BT = min(64, triton.next_power_of_2( + triton.cdiv(max(16, B * T), get_multiprocessor_count(grad_out.device.index)), + )) + if cu_seqlens is not None: + chunk_idx = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) + NT = len(chunk_idx) + else: + chunk_idx = None + NT = triton.cdiv(T, BT) + + BD = triton.next_power_of_2(D) + + def grid(meta): + return (triton.cdiv(D, meta['BD']), NT, N) + + k_update_bwd_kernel_long[grid]( + grad_out, k, a, ka, + dk, da, dka_tmp, + cu_seqlens, chunk_idx, + T, D, + BD=BD, BT=BT, + ) + + if dka_tmp.dim() == 3: + dka = dka_tmp.sum(dim=(0, 1), keepdim=True).type_as(ka) + else: + dka = dka_tmp.sum(dim=(0, 1)).type_as(ka) + + return dk, da, dka + + +class KUpdateFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward(ctx, k, a, ka, cu_seqlens=None, cu_seqlens_cpu=None): + out, use_short, N, T = k_update_fwd(k, a, ka, cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu) + ctx.save_for_backward(k, a, ka) + ctx.use_short = use_short + ctx.N = N + ctx.T = T + ctx.cu_seqlens = cu_seqlens + ctx.cu_seqlens_cpu = cu_seqlens_cpu + return out + + @staticmethod + @input_guard + def backward(ctx, grad_output): + k, a, ka = ctx.saved_tensors + dk, da, dka = k_update_bwd( + grad_output, k, a, ka, + ctx.cu_seqlens, + ctx.use_short, + ctx.N, + ctx.T, + cu_seqlens_cpu=ctx.cu_seqlens_cpu, + ) + return dk, da, dka, None, None + + +def fused_k_rwkv7(k, a, ka, cu_seqlens=None, cu_seqlens_cpu=None): + if k.shape[1] == 1: + return k_update_ref(k, a, ka) + return KUpdateFunction.apply(k, a, ka, cu_seqlens, cu_seqlens_cpu) diff --git a/fla/ops/rwkv7/fused_recurrent.py b/fla/ops/rwkv7/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..2237d4a2b6177efe28342f5068e00c2b381e746f --- /dev/null +++ b/fla/ops/rwkv7/fused_recurrent.py @@ -0,0 +1,333 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.generalized_delta_rule import fused_recurrent_dplr_delta_rule +from fla.ops.utils.op import exp +from fla.utils import USE_CUDA_GRAPH, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BV in [16, 32, 64] + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=['BK'], + use_cuda_graph=USE_CUDA_GRAPH, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_rwkv7_fwd_kernel( + r, + w, + k, + v, + kk, + a, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + REVERSE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_DECODE: tl.constexpr, +): + i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_r = r + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_w = w + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_k = k + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_v = v + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v + p_a = a + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + p_kk = kk + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k + + p_o = o + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + if IS_DECODE: + b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale + b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) + b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32) + b_act_a = -b_kk + b_b = b_kk * b_a + + b_h = exp(b_w)[:, None] * b_h + b_b[:, None] * tl.sum(b_act_a[:, None] * b_h, 0)[None, :] + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_r[:, None], 0) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + else: + for _ in range(0, T): + b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale + b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) + b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32) + b_act_a = -b_kk + b_b = b_kk * b_a + + b_h = exp(b_w)[:, None] * b_h + b_b[:, None] * tl.sum(b_act_a[:, None] * b_h, 0)[None, :] + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_r[:, None], 0) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + p_r += (-1 if REVERSE else 1) * H*K + p_w += (-1 if REVERSE else 1) * H*K + p_k += (-1 if REVERSE else 1) * H*K + p_v += (-1 if REVERSE else 1) * H*V + p_a += (-1 if REVERSE else 1) * H*K + p_kk += (-1 if REVERSE else 1) * H*K + p_o += (-1 if REVERSE else 1) * H*V + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@input_guard +def fused_recurrent_rwkv7_fwd( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + scale: float | None = 1.0, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + IS_DECODE = (T == 1) + + h0 = initial_state + if not output_final_state: + ht = None + else: + ht = r.new_empty(N, H, K, V, dtype=torch.float32) + o = torch.empty_like(v) + + def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) + fused_recurrent_rwkv7_fwd_kernel[grid]( + r, + w, + k, + v, + kk, + a, + o, + h0, + ht, + cu_seqlens, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + REVERSE=reverse, + IS_DECODE=IS_DECODE, + ) + return o, ht + + +def fused_recurrent_rwkv7( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +): + """ + Args: + r (torch.Tensor): + r of shape `[B, T, H, K]`. + w (torch.Tensor): + log decay of shape `[B, T, H, K]`. + k (torch.Tensor): + k of shape `[B, T, H, K]`. + v (torch.Tensor): + v of shape `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, T, H, K]`. + scale (float): + scale of the attention. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (torch.Tensor): + initial state of shape `[B, H, K, V]` if cu_seqlens is None else `[N, H, K, V]` where N = len(cu_seqlens) - 1. + output_final_state (bool): + whether to output the final state. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + elif r.shape[1] < r.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + return fused_recurrent_dplr_delta_rule( + q=r, + k=k, + v=v, + a=a, + b=b, + gk=w, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + + +def fused_mul_recurrent_rwkv7( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + scale: float | None = 1.0, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. + + Args: + r (torch.Tensor): + queries of shape `[B, T, H, K]`. + w (torch.Tensor): + keys of shape `[B, T, H, K]`. + k (torch.Tensor): + values of shape `[B, T, H, V]`. + v (torch.Tensor): + a of shape `[B, T, H, K]`. + kk (torch.Tensor): + b of shape `[B, T, H, K]`. + a (torch.Tensor): + gk of shape `[B, T, H, K]`. decay term in log space! + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: 1. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (Optional[torch.Tensor]): + Cumulative sequence lengths of shape `[N + 1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + elif r.shape[1] < r.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if r.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = r.shape[-1] ** -0.5 + o, final_state = fused_recurrent_rwkv7_fwd( + r, + w, + k, + v, + kk, + a, + scale, + initial_state, + output_final_state, + reverse, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/rwkv7/gate_output_correction.py b/fla/ops/rwkv7/gate_output_correction.py new file mode 100644 index 0000000000000000000000000000000000000000..151071c81cb71198e353c45036d75eb4dd7c9b82 --- /dev/null +++ b/fla/ops/rwkv7/gate_output_correction.py @@ -0,0 +1,245 @@ + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +def gate_output_correction_ref( + o: torch.Tensor, + r: torch.Tensor, + k: torch.Tensor, + r_k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, +): + """ + o: [B, T, H*D] + r: [B, T, H, D] + k: [B, T, H, D] + r_k: [H, D] + v: [B, T, H, D] + g: [B, T, H*D] + """ + # Unsqueeze r_k for broadcasting across batch and time + correction_term = ((r * k * r_k.unsqueeze(0).unsqueeze(0)).sum(-1, keepdim=True) * v).view(o.shape) + output = (o + correction_term) * g + return output + + +def gate_output_correction_backward_ref(grad_output, o, r, k, r_k, v, g): + """ + Reference backward pass implementation in pure PyTorch. + """ + B, T, HD = o.shape + H, D = r.shape[-2], r.shape[-1] + + # Unsqueeze r_k for broadcasting + r_k_b = r_k.unsqueeze(0).unsqueeze(0) + correction_scalar = (r * k * r_k_b).sum(-1, keepdim=True) + gated_input = o + (correction_scalar * v).view(B, T, HD) + + grad_g = grad_output * gated_input + grad_gated_input = grad_output * g + grad_o = grad_gated_input + grad_correction = grad_gated_input + grad_correction_reshaped = grad_correction.view(B, T, H, D) + grad_v = grad_correction_reshaped * correction_scalar + grad_correction_scalar = (grad_correction_reshaped * v).sum(-1, keepdim=True) + grad_r_mul_k_mul_rk = grad_correction_scalar.expand_as(r) + grad_r = grad_r_mul_k_mul_rk * k * r_k_b + grad_k = grad_r_mul_k_mul_rk * r * r_k_b + # Sum over batch and time, keep the head dimension + grad_r_k = (grad_r_mul_k_mul_rk * r * k).sum(dim=(0, 1)) + return grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for num_warps in [2, 4, 8] + for BT in [2, 4, 8] + ], + key=['num_heads', 'head_dim', 'BLOCK_SIZE_D'], + **autotune_cache_kwargs, +) +@triton.jit +def gate_output_correction_fwd_kernel( + o_ptr, r_ptr, k_ptr, r_k_ptr, v_ptr, g_ptr, output_ptr, + o_b_stride, o_t_stride, + r_b_stride, r_t_stride, r_h_stride, + v_b_stride, v_t_stride, v_h_stride, + r_k_h_stride, + T, + T_OFFSET, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + BT: tl.constexpr, +): + pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) + pid_h = tl.program_id(2) + t_start = pid_t_block * BT + T_OFFSET + t_idx = t_start + tl.arange(0, BT)[:, None] + mask_t = t_idx < T + + d_idx = tl.arange(0, BLOCK_SIZE_D)[None, :] + mask_d = d_idx < head_dim + mask = mask_t & mask_d + + offset_rk_h = pid_h * r_k_h_stride + vec_r_k = tl.load(r_k_ptr + offset_rk_h + d_idx, mask=mask_d, other=0.0).to(tl.float32) + + offset_rh = pid_b * r_b_stride + t_idx * r_t_stride + pid_h * r_h_stride + offset_vh = pid_b * v_b_stride + t_idx * v_t_stride + pid_h * v_h_stride + vec_r = tl.load(r_ptr + offset_rh + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_k = tl.load(k_ptr + offset_rh + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_v = tl.load(v_ptr + offset_vh + d_idx, mask=mask, other=0.0).to(tl.float32) + correction = tl.sum(vec_r * vec_k * vec_r_k, axis=1)[:, None] * vec_v + + offset_o = pid_b * o_b_stride + t_idx * o_t_stride + pid_h * head_dim + vec_o = tl.load(o_ptr + offset_o + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_g = tl.load(g_ptr + offset_o + d_idx, mask=mask, other=0.0).to(tl.float32) + final_output = (vec_o + correction) * vec_g + + tl.store(output_ptr + offset_o + d_idx, final_output.to(output_ptr.dtype.element_ty), mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [1, 2, 4] + for BT in [2, 4, 8] + ], + key=['num_heads', 'head_dim', 'BLOCK_SIZE_D'], + **autotune_cache_kwargs, +) +@triton.jit +def gate_output_correction_bwd_kernel( + grad_output_ptr, o_ptr, r_ptr, k_ptr, r_k_ptr, v_ptr, g_ptr, + grad_o_ptr, grad_r_ptr, grad_k_ptr, grad_r_k_intermediate_ptr, grad_v_ptr, grad_g_ptr, + r_b_stride, r_t_stride, r_h_stride, o_b_stride, o_t_stride, r_k_h_stride, + T, + T_OFFSET, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + BT: tl.constexpr, +): + pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) + pid_h = tl.program_id(2) + + t_idx = pid_t_block * BT + T_OFFSET + tl.arange(0, BT)[:, None] + mask_t = t_idx < T + + d_idx = tl.arange(0, BLOCK_SIZE_D)[None, :] + mask_d = d_idx < head_dim + mask = mask_t & mask_d + + rkv_offset = pid_b * r_b_stride + t_idx * r_t_stride + pid_h * r_h_stride + vec_r = tl.load(r_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_k = tl.load(k_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_v = tl.load(v_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + + og_offset = pid_b * o_b_stride + t_idx * o_t_stride + pid_h * head_dim + vec_o = tl.load(o_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_g = tl.load(g_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + vec_grad_output = tl.load(grad_output_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) + + offset_rk_h = pid_h * r_k_h_stride + vec_r_k = tl.load(r_k_ptr + offset_rk_h + d_idx, mask=mask_d, other=0.0).to(tl.float32) + + prod_r_k_rk = vec_r * vec_k * vec_r_k + corr_scalar = tl.sum(prod_r_k_rk, axis=1) + corr_vec = corr_scalar[:, None] * vec_v + gated_input = vec_o + corr_vec + + vec_grad_g = vec_grad_output * gated_input + vec_grad_gate = vec_grad_output * vec_g + vec_grad_o = vec_grad_gate + vec_grad_corr = vec_grad_gate + vec_grad_v = vec_grad_corr * corr_scalar[:, None] + grad_corr_s = tl.sum(vec_grad_corr * vec_v, axis=1)[:, None] + vec_grad_r = grad_corr_s * vec_k * vec_r_k + vec_grad_k = grad_corr_s * vec_r * vec_r_k + local_grad_rk = grad_corr_s * vec_r * vec_k + + tl.store(grad_o_ptr + og_offset + d_idx, vec_grad_o.to(grad_o_ptr.dtype.element_ty), mask=mask) + tl.store(grad_g_ptr + og_offset + d_idx, vec_grad_g.to(grad_g_ptr.dtype.element_ty), mask=mask) + tl.store(grad_r_ptr + rkv_offset + d_idx, vec_grad_r.to(grad_r_ptr.dtype.element_ty), mask=mask) + tl.store(grad_k_ptr + rkv_offset + d_idx, vec_grad_k.to(grad_k_ptr.dtype.element_ty), mask=mask) + tl.store(grad_v_ptr + rkv_offset + d_idx, vec_grad_v.to(grad_v_ptr.dtype.element_ty), mask=mask) + tl.store(grad_r_k_intermediate_ptr + rkv_offset + d_idx, + local_grad_rk.to(grad_r_k_intermediate_ptr.dtype.element_ty), mask=mask) + + +def gate_output_correction_backward_triton(grad_output, o, r, k, r_k, v, g): + batch_size, seq_len, _ = o.shape + num_heads, head_dim = r.shape[-2], r.shape[-1] + + grad_o = torch.empty_like(o) + grad_r = torch.empty_like(r) + grad_k = torch.empty_like(k) + grad_v = torch.empty_like(v) + grad_g = torch.empty_like(g) + # Keep intermediate in float32 for precision + grad_r_k = torch.empty_like(r, dtype=torch.float32) + + BLOCK_SIZE_D = triton.next_power_of_2(head_dim) + + for t_offset in range(0, seq_len, 65536): + T_SIZE = min(65536, seq_len - t_offset) + def grid(meta): return (batch_size, triton.cdiv(T_SIZE, meta['BT']), num_heads) + + gate_output_correction_bwd_kernel[grid]( + grad_output, o, r, k, r_k, v, g, + grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g, + r.stride(0), r.stride(1), r.stride(2), + o.stride(0), o.stride(1), + r_k.stride(0), + T_SIZE, t_offset, + num_heads=num_heads, head_dim=head_dim, BLOCK_SIZE_D=BLOCK_SIZE_D, + ) + # Sum over batch and time to get the final gradient for r_k + grad_r_k = grad_r_k.sum(dim=(0, 1)).type_as(r_k) + return grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g + + +class GateOutputCorrection(torch.autograd.Function): + @staticmethod + @autocast_custom_fwd + @input_guard + def forward(ctx, o, r, k, r_k, v, g): + assert r_k.dim() == 2 and r_k.shape[0] == r.shape[-2] and r_k.shape[1] == r.shape[-1] + + batch_size, seq_len, _ = o.shape + num_heads, head_dim = r.shape[-2], r.shape[-1] + output = torch.empty_like(o) + ctx.save_for_backward(o, r, k, r_k, v, g) + for t in range(0, seq_len, 65536): + T_OFFSET = t + T_SIZE = min(65536, seq_len - t) + def grid(meta): return (batch_size, triton.cdiv(T_SIZE, meta['BT']), num_heads) + + gate_output_correction_fwd_kernel[grid]( + o, r, k, r_k, v, g, output, + o.stride(0), o.stride(1), + r.stride(0), r.stride(1), r.stride(2), + v.stride(0), v.stride(1), v.stride(2), + r_k.stride(0), + T_SIZE, T_OFFSET, + num_heads, head_dim, BLOCK_SIZE_D=triton.next_power_of_2(head_dim), + ) + return output + + @staticmethod + @autocast_custom_bwd + @input_guard + def backward(ctx, grad_output): + o, r, k, r_k, v, g = ctx.saved_tensors + return gate_output_correction_backward_triton(grad_output, o, r, k, r_k, v, g) + + +gate_output_correction = GateOutputCorrection.apply diff --git a/fla/ops/simple_gla/README.md b/fla/ops/simple_gla/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c359ced5ed1304fdb6bf3edb76cc37470064abf0 --- /dev/null +++ b/fla/ops/simple_gla/README.md @@ -0,0 +1,10 @@ +# Simple GLA + +Gating mechanism in [Gated RFA](https://arxiv.org/abs/2103.02143), [Mamba2](https://arxiv.org/abs/2405.21060) and [YOCO](https://arxiv.org/abs/2405.05254) (a.k.a., Gated RetNet). + +Compared to GLA, the gating is head-wise instead of elementwise. +As a result, we can adapt the RetNet kernel for training using matmul w/o numerical instability. +It is faster than GLA but has less expressive power. +I will use it as a baseline for the GLA. + +$S_{t+1} = g_{t+1} \odot S_{t} + K_{t+1} V_{t+1}^{\top}$ where $g$ is a scalar. diff --git a/fla/ops/simple_gla/__init__.py b/fla/ops/simple_gla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a0e62619139ef18b7ef030354d5347c4b4062a37 --- /dev/null +++ b/fla/ops/simple_gla/__init__.py @@ -0,0 +1,12 @@ + +from .chunk import chunk_simple_gla +from .fused_chunk import fused_chunk_simple_gla +from .fused_recurrent import fused_recurrent_simple_gla +from .parallel import parallel_simple_gla + +__all__ = [ + 'chunk_simple_gla', + 'fused_chunk_simple_gla', + 'fused_recurrent_simple_gla', + 'parallel_simple_gla', +] diff --git a/fla/ops/simple_gla/chunk.py b/fla/ops/simple_gla/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..9cb5191cb7d00f04003794e0b53ff3d5cdf63f8b --- /dev/null +++ b/fla/ops/simple_gla/chunk.py @@ -0,0 +1,316 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton + +from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv, chunk_fwd_o +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_simple_gla_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + h, ht = chunk_fwd_h( + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=None, + gv=None, + h0=initial_state, + output_final_state=output_final_state, + states_in_fp32=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + o = chunk_fwd_o( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return o, ht + + +def chunk_simple_gla_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_gamma: torch.Tensor, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # (SY 09/22) states_in_fp32 seems not affecting the error of dg but for safety, set to True + h, _ = chunk_fwd_h( + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=None, + gv=None, + h0=initial_state, + output_final_state=False, + states_in_fp32=True, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + gk=None, + gv=None, + do=do, + h0=initial_state, + dht=dht, + scale=scale, + states_in_fp32=True, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + dq, dk, _, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + h=h, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dv = chunk_bwd_dv( + q=q, + k=k, + g=g, + g_gamma=g_gamma, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, dg, dh0 + + +class ChunkSimpleGLAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + g, + g_gamma, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ): + T = q.shape[1] + chunk_size = min(64, max(16, triton.next_power_of_2(T))) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices) if g is not None else None + o, ht = chunk_simple_gla_fwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, k, v, g, g_gamma, initial_state, chunk_indices) + ctx.chunk_size = chunk_size + ctx.scale = scale + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), ht + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht): + chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens + q, k, v, g, g_gamma, initial_state, chunk_indices = ctx.saved_tensors + dq, dk, dv, dg, dh0 = chunk_simple_gla_bwd( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + initial_state=initial_state, + do=do, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + if g is not None: + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices).to(g) + else: + dg = None + return dq.to(q), dk.to(k), dv.to(v), dg, None, None, dh0, None, None, None + + +@torch.compiler.disable +def chunk_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + Compared to GLA, the gating is head-wise instead of elementwise. + g_gamma (torch.Tensor): + Log decay of shape `[H]`. + Head-wise data-independent decay is used if `g_gamma` is provided. + Only one of `g` or `g_gamma` should be provided. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.simple_gla import chunk_simple_gla + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, device='cuda')) + >>> o, ht = chunk_simple_gla( + q, k, v, g, + initial_state=None, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = chunk_simple_gla( + q, k, v, g, + initial_state=None, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkSimpleGLAFunction.apply( + q, + k, + v, + g, + g_gamma, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, final_state diff --git a/fla/ops/simple_gla/fused_chunk.py b/fla/ops/simple_gla/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..21d8f9c10ce83ffa1713e0610c91e5c7659c3cc3 --- /dev/null +++ b/fla/ops/simple_gla/fused_chunk.py @@ -0,0 +1,106 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.common.fused_chunk import fused_chunk + + +def fused_chunk_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor = None, + g_gamma: torch.Tensor = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + Compared to GLA, the gating is head-wise instead of elementwise. + g_gamma (torch.Tensor): + Log decay of shape `[H]`. + Head-wise data-independent decay is used if `g_gamma` is provided. + Only one of `g` or `g_gamma` should be provided. + scale (Optional[int]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.simple_gla import fused_chunk_simple_gla + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_chunk_simple_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = fused_chunk_simple_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = fused_chunk( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/simple_gla/fused_recurrent.py b/fla/ops/simple_gla/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..ebdd8807d432ccfb4e0e5d55d5af7f25b13453c3 --- /dev/null +++ b/fla/ops/simple_gla/fused_recurrent.py @@ -0,0 +1,110 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch + +from fla.ops.common.fused_recurrent import fused_recurrent + + +def fused_recurrent_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor = None, + g_gamma: torch.Tensor = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + reverse: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + Compared to GLA, the gating is head-wise instead of elementwise. + g_gamma (torch.Tensor): + Log decay of shape `[H]`. + Head-wise data-independent decay is used if `g_gamma` is provided. + Only one of `g` or `g_gamma` should be provided. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.simple_gla import fused_recurrent_simple_gla + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_recurrent_simple_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_simple_gla( + q, k, v, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = fused_recurrent( + q=q, + k=k, + v=v, + g=g, + g_gamma=g_gamma, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/simple_gla/naive.py b/fla/ops/simple_gla/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..ab452f590e2d6bf443dfe2dba479e037b0863208 --- /dev/null +++ b/fla/ops/simple_gla/naive.py @@ -0,0 +1,112 @@ + + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_chunk_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + scale: float | None = None, +): + q, k, v, g = map(lambda x: rearrange(x, 'b t h ... -> b h t ...').to(torch.float32), [q, k, v, g]) + if scale is None: + scale = 1.0 / q.shape[-1] ** 0.5 + + T = q.shape[-2] + BT = chunk_size + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + # Pad all tensors + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + g = F.pad(g, (0, pad_len)) + decay = g + B, H, T1, K = q.shape + V = v.shape[-1] + q = q * scale + q, k, v, decay = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, decay.unsqueeze(-1)]) + decay = decay.squeeze(-1).cumsum(-1) + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + S = k.new_zeros(B, H, K, V) + if initial_state is not None: + S = initial_state + o = torch.zeros_like(v) + for i in range(0, T1 // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]) + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_i + S = S * decay[:, :, i, -1, None, None].exp() + \ + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_i + if not output_final_state: + S = None + # unpad + o = rearrange(o, 'b h n c d -> b (n c) h d')[:, :T] + return o, S + + +def naive_recurrent_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = True, +): + dtype = q.dtype + q, k, v, g = map(lambda x: x.transpose(1, 2).float(), (q, k, v, g)) + B, H, T, K = q.shape + V = v.shape[-1] + if scale is None: + scale = K ** -0.5 + q = q * scale + o = v.new_zeros(B, H, T, V) + + S = q.new_zeros(B, H, K, V) + if initial_state is not None: + S += initial_state + + for i in range(T): + gate = g[:, :, i].exp() + key = k[:, :, i] + value = v[:, :, i] + kv = key.unsqueeze(-1) * value.unsqueeze(-2) + S = S * gate.unsqueeze(-1).unsqueeze(-1) + kv + q_i = q[:, :, i, :] + o_i = (q_i.unsqueeze(-1) * S).sum(-2) + o[:, :, i] = o_i + if not output_final_state: + S = None + return o.transpose(1, 2).to(dtype), S + + +def naive_parallel_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, +): + q, k, v, g = map(lambda x: rearrange(x, 'b t h ... -> b h t ...').to(torch.float32), [q, k, v, g]) + if scale is None: + scale = 1.0 / q.shape[-1] ** 0.5 + dtype = q.dtype + A = (q @ k.transpose(-1, -2) * scale) + if g is not None: + g = g.cumsum(-1) + D = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().tril() + A = A * D + else: + A = A.tril() + o = A @ v + o = o.transpose(1, 2) + return o.to(dtype), A diff --git a/fla/ops/simple_gla/parallel.py b/fla/ops/simple_gla/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..ad4e71daee89772fcdc33de93f94ccad25f95367 --- /dev/null +++ b/fla/ops/simple_gla/parallel.py @@ -0,0 +1,734 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cumsum import chunk_global_cumsum, chunk_local_cumsum +from fla.ops.utils.op import exp +from fla.utils import ( + IS_INTEL_ALCHEMIST, + IS_NVIDIA_HOPPER, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_shared_mem, + input_guard, +) + +# https://github.com/intel/intel-xpu-backend-for-triton/issues/3449 +triton_config = {'grf_mode': 'large'} if IS_INTEL_ALCHEMIST else {} +NUM_WARPS = [2, 4, 8] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + + +@triton.heuristics({ + 'NV': lambda args: triton.cdiv(args['V'], args['BV']), + 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None, + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8, 16] + for num_stages in [2, 3, 4] + ], + key=["BT", "BS", "BK", "BV", "USE_G"], + **autotune_cache_kwargs, +) +@triton.jit +def parallel_simple_gla_fwd_kernel( + q, + k, + v, + g, + o, + attn, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + OUTPUT_ATTENTIONS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v = i_kv // NV, i_kv % NV + i_b, i_h = i_bh // H, i_bh % H + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + o += ((i_k * all + bos) * H + i_h) * V + if USE_G: + g += bos * H + i_h + if OUTPUT_ATTENTIONS: + attn += i_k * B * H * T * T + (bos * H + i_h * T) * T + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + b_o = tl.zeros([BT, BV], dtype=tl.float32) + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + m_q = o_q < T + # Q block and K block have overlap. + # masks required + if USE_G: + # [BT,] + b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')).to(tl.float32) + # rescale interchunk output + else: + b_gq = None + + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) + b_s = tl.dot(b_q, b_k) + if USE_G: + b_gk = tl.load(g + o_k * H, mask=m_k, other=0) + b_s *= exp(b_gq[:, None] - b_gk[None, :]) + b_s = tl.where(m_s, b_s, 0) + # [BT, BV] + if i_s >= 0: + b_o += tl.dot(b_s.to(b_q.dtype), b_v) + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + for i_s in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BS] + m_s = m_q[:, None] & m_k[None, :] + b_s = tl.dot(b_q, b_k) + if USE_G: + b_g = tl.load(g + o_k * H, mask=m_k, other=0) + b_gn = tl.load(g + (min(i_s + BS, T) - 1) * H) + b_gp = tl.load(g + (i_s-1) * H) if i_s % BT > 0 else 0. + # No concrete meaning. Just to avoid some layout bugs. + b_s *= exp(b_gq[:, None] + (b_gn - b_g)[None, :]) + b_gq += b_gn - b_gp + b_s = tl.where(m_s, b_s, 0) + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + if i_s >= 0: + b_o += tl.dot(b_s.to(b_v.dtype), b_v) + p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def parallel_simple_gla_bwd_kernel_dq( + i_t, + i_k, + i_v, + q, + k, + v, + g, + do, + dq, + dg, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, +): + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT, BK] + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + m_q = o_q < T + for i_s in range(0, i_t * BT, BS): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BV] @ [BV, BS] = [BT, BS] + b_ds = tl.dot(b_do, b_v) + if USE_G: + b_g = tl.load(g + o_k * H, mask=m_k, other=0) + b_gn = tl.load(g + (min(i_s + BS, T) - 1) * H) + b_gp = tl.load(g + (i_s - 1) * H) if i_s % BT > 0 else 0. + b_ds *= tl.where(m_k, exp(b_gn - b_g), 0)[None, :] + if i_s > 0: + b_dq *= exp(b_gn - b_gp) + # [BT, BS] @ [BS, BK] = [BT, BK] + b_dq += tl.dot(b_ds.to(b_v.dtype), b_k) + + if USE_G: + # [BT,] + b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) + # [BT, BK] + b_dq *= exp(b_gq)[:, None] + + # Q block and K block have overlap. masks required + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + # [BS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BV] @ [BV, BS] = [BT, BS] + b_ds = tl.dot(b_do, b_v) + if USE_G: + b_gk = tl.load(g + o_k * H, mask=m_k, other=0) + b_ds *= exp(b_gq[:, None] - b_gk[None, :]) + m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) + b_ds = tl.where(m_s, b_ds, 0) + # [BT, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + + b_dq *= scale + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + if USE_G: + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_dg = tl.sum(b_dq * b_q, 1) + p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@triton.jit(do_not_specialize=['T']) +def parallel_simple_gla_bwd_kernel_dkv( + i_t, + i_k, + i_v, + q, + k, + v, + g, + do, + dk, + dv, + dg, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, +): + o_k = i_t * BT + tl.arange(0, BT) + m_k = o_k < T + # [BT, BK] + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + # [BT, BV] + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + if USE_G: + b_gk = tl.load(g + o_k * H, mask=m_k, other=0) + NTS = tl.cdiv(T, BS) + # [BT, BK] + for i_s in range(NTS * BS - BS, (i_t + 1) * BT - BS, -BS): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + + o_q = i_s + tl.arange(0, BS) + m_q = o_q < T + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT, BS] + b_ds = tl.dot(b_v, tl.trans(b_do)) + b_s = tl.dot(b_k, tl.trans(b_q)) + if USE_G: + b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) + b_gp = tl.load(g + (min(i_s + BS, T) - 1) * H) + b_gn = tl.load(g + (i_s - 1) * H) if i_s % BT > 0 else 0. + if i_s >= 0: + b_gpn = exp(b_gp - b_gn) + b_dk *= b_gpn + b_dv *= b_gpn + b_gqn = exp(b_gq - b_gn) + b_ds *= b_gqn[None, :] + b_s *= b_gqn[None, :] + # [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + # [BT, BV] + b_dv += tl.dot(b_s.to(b_do.dtype), b_do) + + if USE_G: + b_gn = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) + if i_t >= 0: + b_gpn = exp(b_gn - b_gk)[:, None] + b_dk *= b_gpn + b_dv *= b_gpn + + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + + o_q = i_s + tl.arange(0, BS) + m_q = o_q < T + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BS] + b_s = tl.dot(b_k, tl.trans(b_q)) + b_ds = tl.dot(b_v, tl.trans(b_do)) + if USE_G: + b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) + if i_s >= 0: + b_gkq = exp(-b_gk[:, None] + b_gq[None, :]) + b_ds *= b_gkq + b_s *= b_gkq + m_s = o_k[:, None] <= o_q[None, :] + b_s = tl.where(m_s, b_s, 0) + b_ds = tl.where(m_s, b_ds, 0) + # [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + b_dv += tl.dot(b_s.to(b_do.dtype), b_do) + b_dk *= scale + b_dv *= scale + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + if USE_G: + b_dg = tl.load(dg + o_k * H, mask=m_k, other=0) + b_dg -= tl.sum(b_dk * b_k, 1) + tl.store(dg + o_k * H, b_dg.to(dg.dtype.element_ty), mask=m_k) + + +@triton.heuristics({ + 'NV': lambda args: triton.cdiv(args['V'], args['BV']), + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config(triton_config, num_warps=num_warps) + for num_warps in NUM_WARPS + ], + key=['BT', 'BS', 'BK', 'BV', 'USE_G'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def parallel_simple_gla_bwd_kernel( + q, + k, + v, + g, + do, + dq, + dk, + dv, + dg, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v = i_kv // NV, i_kv % NV + i_b, i_h = i_bh // H, i_bh % H + dq += i_v * B * H * T * K + dk += i_v * B * H * T * K + dv += i_k * B * H * T * V + if USE_G: + dg += i_kv * B * H * T + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + if USE_G: + g += bos * H + i_h + dg += bos * H + i_h + + parallel_simple_gla_bwd_kernel_dq( + i_t=i_t, + i_k=i_k, + i_v=i_v, + q=q, + k=k, + v=v, + g=g, + do=do, + dq=dq, + dg=dg, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + USE_G=USE_G, + ) + tl.debug_barrier() + parallel_simple_gla_bwd_kernel_dkv( + i_t=i_t, + i_k=i_k, + i_v=i_v, + q=q, + k=k, + v=v, + g=g, + do=do, + dk=dk, + dv=dv, + dg=dg, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + USE_G=USE_G, + ) + + +def parallel_simple_gla_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float, + output_attentions: bool = False, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT, BS = chunk_size, 32 + if check_shared_mem('hopper', k.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + elif check_shared_mem('ampere', k.device.index): + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + else: + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert BT % BS == 0 + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + # local cumulative decay in log space + if g is not None: + g = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + grid = (NK * NV, NT, B * H) + o = torch.empty(NK, *v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) + attn = q.new_zeros(NK, B, H, T, T) if output_attentions else None + + parallel_simple_gla_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + o=o, + attn=attn, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + H=H, + T=T, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + ) + o = o.sum(0) + + if output_attentions: + attn = attn.sum(0) + return o, g, attn + + +def parallel_simple_gla_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + do: torch.Tensor, + scale: float, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT, BS = chunk_size, 32 + if check_shared_mem('hopper', k.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + elif check_shared_mem('ampere', k.device.index): + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + elif check_shared_mem('ada', k.device.index): + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + else: + BK = min(32, triton.next_power_of_2(K)) + BV = min(32, triton.next_power_of_2(V)) + + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert BT % BS == 0 + + dq = torch.empty(NV, * q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) + dk = torch.empty(NV, * k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) + dv = torch.empty(NK, * v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) + dg = torch.empty(NK*NV, *g.shape, dtype=torch.float, device=q.device) if g is not None else None + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + grid = (NK * NV, NT, B * H) + parallel_simple_gla_bwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + do=do, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + dg = chunk_global_cumsum(dg.sum(0), reverse=True, cu_seqlens=cu_seqlens) if g is not None else None + return dq, dk, dv, dg + + +class ParallelSimpleGLAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, g, scale, output_attentions, cu_seqlens, cu_seqlens_cpu): + chunk_size = 128 + ctx.dtype = q.dtype + + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + o, g, attn = parallel_simple_gla_fwd( + q=q, + k=k, + v=v, + g=g, + scale=scale, + output_attentions=output_attentions, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, k, v, g, cu_seqlens, chunk_indices) + ctx.scale = scale + ctx.chunk_size = chunk_size + return o.to(q.dtype), attn + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, da=None): + q, k, v, g, cu_seqlens, chunk_indices = ctx.saved_tensors + dq, dk, dv, dg = parallel_simple_gla_bwd( + q=q, + k=k, + v=v, + g=g, + do=do, + scale=ctx.scale, + chunk_size=ctx.chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), dg.to(ctx.dtype) if dg is not None else None, None, None, None, None + + +def parallel_simple_gla( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + scale: float | None = None, + output_attentions: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + Forget gates of shape `[B, T, H]`. + Compared to GLA, the gating is head-wise instead of elementwise. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + output_attentions (bool): + Whether to output the materialized attention scores of shape [B, H, T, T]. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + attn (torch.Tensor): + Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None` + """ + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if output_attentions: + assert cu_seqlens is None, "output_attentions=True is not supported with variable-length sequences" + + if scale is None: + scale = k.shape[-1] ** -0.5 + o, attn = ParallelSimpleGLAFunction.apply( + q, + k, + v, + g, + scale, + output_attentions, + cu_seqlens, + cu_seqlens_cpu, + ) + return o, attn diff --git a/fla/ops/titans/__init__.py b/fla/ops/titans/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd3f45dee4123a33bed73f04edf170ec70a7bd71 --- /dev/null +++ b/fla/ops/titans/__init__.py @@ -0,0 +1,6 @@ + +from .naive import chunk_titans_linear + +__all__ = [ + 'chunk_titans_linear', +] diff --git a/fla/ops/titans/log_impl.py b/fla/ops/titans/log_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..49e6a90bb29058ae7f4924cc23891cd3abd4f4ee --- /dev/null +++ b/fla/ops/titans/log_impl.py @@ -0,0 +1,153 @@ +import torch + + +def cal_n_log(log_theta, log_eta, seq_len): + """ + calculate n_{i,j} in log space + log(n_{i,j}) = log(θ_j) + sum_{k=j+1}^i log(η_k) + """ + # create log(n) + log_n = torch.zeros(*log_theta.shape, seq_len, dtype=log_eta.dtype).to( + log_eta.device, + ) # [batch_size, num_heads, seq_len, seq_len] + for i in range(seq_len): + for j in range(i + 1): + if i == j: + log_n[..., j, i] = log_theta[..., j] + else: + log_n[..., j, i] = log_theta[..., j] + torch.sum( + log_eta[..., j + 1: i + 1], dim=-1, + ) + + return log_n + + +def cal_f_log(log_beta, seq_len, log_m): + """ + cal_f_log(log_beta, seq_len, log_m) -> f + log(f_t) = log(sum_{i=1}^t exp(sum_{k=i+1}^t log(1-α_k) + sum_{k=1}^i log(η_k))) + """ + # create f + # f = torch.zeros_like(log_beta) + # for t in range(seq_len): + # for i in range(t + 1): + # f[..., t] += torch.exp(log_beta[..., t] - log_beta[..., i] + log_m[..., i]) + log_f = torch.zeros_like(log_beta) + for t in range(seq_len): + a_i = log_beta[..., t: t + 1] - log_beta[..., : t + 1] + log_m[..., : t + 1] + log_f[..., t] = torch.logsumexp(a_i, dim=-1) + f = torch.exp(log_f) + + # this version overflow and even slower + # t_indices = torch.arange(seq_len, device=log_beta.device) + # i_indices = torch.arange(seq_len, device=log_beta.device) + # + # mask = i_indices.unsqueeze(0) <= t_indices.unsqueeze(1) + # log_beta_t = log_beta.unsqueeze(-1) # [..., seq_len, 1] + # log_beta_i = log_beta.unsqueeze(-2) # [..., 1, seq_len] + # log_m_i = log_m.unsqueeze(-2) + # a_i = log_beta_t - log_beta_i + log_m_i + # masked_a_i = torch.where(mask, a_i, torch.tensor(-float('inf'), device=a_i.device, dtype=a_i.dtype)) + # log_f = torch.logsumexp(masked_a_i, dim=-1) # [..., seq_len] + # + # f = torch.exp(log_f) + return f + + +def cal_G_log(log_beta, log_n, seq_len): + """ + calculate G_{i,j} + log(G_{i,j}) = log(sum_{k=j}^i exp(log(β_i/β_k) + log(n_{k,j}))) + """ + # G = torch.zeros(*log_beta.shape[:-1], seq_len, seq_len, device = log_beta.device) + # # Fill in the lower triangular part + # for i in range(seq_len): # row + # for j in range(i + 1): # column + # # Sum from k=j to i + # for k in range(j, i + 1): + # G[..., i, j] += torch.exp(log_beta[..., i] - log_beta[..., k] + log_n[..., j, k]) + + log_G = torch.full( + (*log_beta.shape[:-1], seq_len, seq_len), float("-inf"), device=log_beta.device, + ) + # fill in the lower triangular part + for i in range(seq_len): # row + for j in range(i + 1): # column + terms = ( + log_beta[..., i: i + 1] + - log_beta[..., j: i + 1] + + log_n[..., j: j + 1, j: i + 1].squeeze(-2) + ) + # use logsumexp to avoid overflow + log_G[..., i, j] = torch.logsumexp(terms, dim=-1) + + G = torch.exp(log_G) + return G + + +def _combine_params_log(log_theta, log_alpha_complement, log_eta, seq_len): + """ + Update rule for Titans in log space + + Parameters: + - log_theta: log(θ) + - log_alpha_complement: log(1-α) + - log_eta: log(η) + - seq_len: sequence length + + Returns: + - log_beta, beta_T, log_f, f_T, log_g, log_G, m_T, n_T + """ + # calculate log(β_t) = sum_{k=1}^t log(1-α_k) + log_beta = torch.cumsum(log_alpha_complement, dim=-1) + + # get β_T + beta_T = torch.exp(log_beta[..., -1]) + + # calculate log(m_i) = sum_{k=1}^i log(η_k) + log_m = torch.cumsum(log_eta, dim=-1) + m_T = torch.exp(log_m[..., -1]) + + # cal log(n_{i,j}) + log_n = cal_n_log(log_theta, log_eta, seq_len) + n_T = torch.exp(log_n[..., -1]) + + # cal log(f_t) + f = cal_f_log(log_beta, seq_len, log_m) + f_T = f[..., -1] + + # cal log(G_{i,j}) + G = cal_G_log(log_beta, log_n, seq_len) + # get log(g_j) = log(G_{T,j}) + g = G[..., -1, :] + + return log_beta, beta_T, f, f_T, g, G, m_T, n_T + + +def combine_params_log(theta, alpha, eta, seq_len): + """ + log space Titians + + Parameters: + - theta: θ + - alpha: α + - eta: η + - seq_len: sequence length + + Returns: + - beta, beta_T, f, f_T, g, G, m_T, n_T + """ + # convert to log space + log_theta = torch.log(theta.squeeze(-1)) + log_alpha_complement = torch.log(1 - alpha.squeeze(-1)) + log_eta = torch.log(eta.squeeze(-1)) + + # combine params in log space + log_beta, beta_T, f, f_T, g, G, m_T, n_T = _combine_params_log( + log_theta, log_alpha_complement, log_eta, seq_len, + ) + + # convert back to normal space + beta = torch.exp(log_beta) + + return beta, beta_T, f, f_T, g, G, m_T, n_T diff --git a/fla/ops/titans/naive.py b/fla/ops/titans/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..39875d4ad680baf4c4a2cd5af22df875cf5a2e10 --- /dev/null +++ b/fla/ops/titans/naive.py @@ -0,0 +1,374 @@ + +import torch +import torch.nn.functional as F + +from fla.ops.titans.log_impl import combine_params_log + + +def cal_n(theta, eta, seq_len): + n = torch.zeros(*theta.shape, seq_len, dtype=theta.dtype).to( + theta.device, + ) # [batch_size, num_heads, seq_len, seq_len] + + # 1. deal with diagonal elements + indices = torch.arange(seq_len, device=theta.device) + n[..., indices, indices] = theta[..., indices] + + # 2. Create a cumulative product matrix + # First create a mask to mark the positions where eta needs to be multiplied + mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).to(theta.device) + # Convert mask to boolean type + mask = mask.bool() + # Expand eta to match the target shape + eta_expanded = eta.unsqueeze(-2).expand(*theta.shape[:-1], seq_len, seq_len) + # Create a matrix filled with 1s for cumulative product + cumulative = torch.ones_like(eta_expanded) + cumulative = torch.where(mask, eta_expanded, cumulative) + # Calculate the cumulative product + cumulative_prod = torch.cumprod(cumulative, dim=-1) + + # 3. Calculate non-diagonal elements + # Create an expanded version of theta + theta_expanded = theta.unsqueeze(-1).expand(*theta.shape[:-1], seq_len, seq_len) + # Create a mask to keep only the upper triangular part (excluding the diagonal) + upper_triangular = torch.triu(torch.ones_like(n), diagonal=1).bool() + # Combine theta and cumulative product + n = torch.where(upper_triangular, theta_expanded * cumulative_prod, n) + return n + + +def cal_f(beta, seq_len, m): + a = torch.tril(beta.to(torch.float32).unsqueeze(-1).expand(*beta.shape, seq_len), 0) + ratio = (m.to(torch.float32) / beta.to(torch.float32)).unsqueeze(-1) + f = torch.matmul(a, ratio).squeeze(-1) + return f.to(beta.dtype) + + +def cal_G(beta, n, seq_len): + i_indices = torch.arange(seq_len, device=beta.device) + j_indices = torch.arange(seq_len, device=beta.device) + k_indices = torch.arange(seq_len, device=beta.device) + beta_ratio = beta[..., :, None] / beta[..., None, :] # [..., i, k] + + # create mask + k_mask = (k_indices[None, None, :] >= j_indices[None, :, None]) & ( + k_indices[None, None, :] <= i_indices[:, None, None] + ) + + # use mask to filter out invalid values + masked_beta_ratio = beta_ratio[..., :, None, :] * k_mask # [..., i, j, k] + masked_n = n[..., None, :, :] * k_mask # [..., i, j, k] + # calculate G + G = torch.sum(masked_beta_ratio * masked_n, dim=-1) # [..., i, j] + return G + + +def combine_params(theta, alpha, eta, seq_len): + theta = theta.squeeze(-1) + eta = eta.squeeze(-1) + alpha = alpha.squeeze(-1) + beta = torch.cumprod(1 - alpha, dim=-1) # β_t = ∏(1 - α_t) in titans paper + beta_T = beta[..., -1] # β_T + # Calculate m_i = ∏(k=1 to i) η_k + m = torch.cumprod(eta, dim=-1) # [batch_size, num_heads, seq_len] + m_T = m[..., -1] # m_T + # Calculate n_{i,j} + # We need to calculate ∏(k=j+1 to i) η_k for each i,j pair + # # this may be optimized + # n = torch.zeros(*theta.shape, seq_len, dtype = theta.dtype).to( + # theta.device) # [batch_size, num_heads, seq_len, seq_len] + # for i in range(seq_len): + # for j in range(i + 1): + # if i == j: + # n[..., j, i] = theta[..., j] + # else: + # # Calculate product of eta from j+1 to i + # eta_product = torch.prod(eta[..., j + 1:i + 1], dim = -1) + # n[..., j, i] = theta[..., j] * eta_product + + n = cal_n(theta, eta, seq_len) + n_T = n[..., -1] # [batch_size, num_heads, seq_len] + # Calculate f_t = ∑(i=1 to t) (β_t/β_i) m_i + # f = torch.zeros_like(theta) + # for t in range(seq_len): + # for i in range(t + 1): + # f[..., t] += (beta[..., t] / beta[..., i]) * m[..., i] + f = cal_f(beta, seq_len, m) + f_T = f[..., -1] # [batch_size, num_heads, seq_len] + # Calculate g_j = ∑(i=j to t) (β_t/β_i) n_{i,j} + # g = torch.zeros_like(theta) # [batch_size, num_heads, seq_len] + # for j in range(seq_len): + # for i in range(j, seq_len): + # g[..., j] += (beta[..., -1] / beta[..., i]) * n[..., j, i] + # G = torch.zeros(*beta.shape[:-1], seq_len, seq_len, device = beta.device) + # # Fill in the lower triangular part + # for i in range(seq_len): # row + # for j in range(i + 1): # column + # # Sum from k=j to i + # for k in range(j, i + 1): + # G[..., i, j] += (beta[..., i] / beta[..., k]) * n[..., j, k] + G = cal_G(beta, n, seq_len) + g = G[:, :, -1, :] # [batch_size, num_heads, seq_len] + # g2, G2 = compute_g_and_G(beta, n, seq_len) + return beta, beta_T, f, f_T, g, G, m_T, n_T + + +def titans_linear( + q, k, v, w, b, theta, alpha, eta, eps, chunk_size, initial_state, output_final_state, +): + """ + Implementation of Titans Linear function based on the update rules: + M_t = (1 - alpha_t) * M_{t-1} + S_t + S_t = eta_t * S_{t-1} - theta_t * nabla_l(M_{t-1}; x_t) + + Args: + q: Query tensor + k: Key tensor + v: Value tensor + w: Weight tensor + b: Bias tensor + theta: Learning rate tensor + alpha: Momentum decay tensor + eta: Step size tensor + eps: Epsilon for numerical stability + initial_state: Initial state M_0 + output_final_state: Whether to output the final state + + Returns: + Tuple of (output tensor, final state) + """ + B, H, T, D = q.shape + device = q.device + w = w.reshape(H, 1, D).to(torch.float32) + b = b.reshape(H, 1, D).to(torch.float32) + # Initialize states + if initial_state is None: + M_prev = torch.zeros(B, H, D, D, device=device) + else: + M_prev = initial_state + M_prev_nabla = M_prev.clone() + S_prev = torch.zeros_like(M_prev) + outputs = [] + + # Process sequence step by step + for t in range(T): + # Get current step inputs + q_t = q[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + k_t = k[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + v_t = v[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + theta_t = theta[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + alpha_t = alpha[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + eta_t = eta[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) + + # Compute gradient + km = k_t @ M_prev_nabla # (batch_size, num_heads, 1, dim) + reconstruction_target = v_t - k_t + mean = km.mean(-1, keepdim=True) + var = km.var(-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + km_hat = (km - mean) / rstd + + grad = w * km_hat + b - reconstruction_target + grad = grad * w + # v_new = (D * grad - grad.sum(-1, keepdim = True) - km_hat * (grad * km_hat).sum(-1, keepdim = True)) / ( + # rstd * D) + v_new = D * grad - grad.sum(-1, keepdim=True) / (rstd * D) + proj_term = km_hat * (grad * km_hat).sum(-1, keepdim=True) / (rstd * D) + v_new = v_new - proj_term + # v_new = grad + + # Update S_t + S_t = eta_t * S_prev - 2 * theta_t * k_t.transpose(-2, -1) @ v_new + + # Update M_t + M_t = (1 - alpha_t) * M_prev + S_t + + # Store output + output_t = q_t @ M_t # (batch_size, num_heads, seq_len, dim) + mean = output_t.mean(dim=-1, keepdim=True) + var = output_t.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + output_t = output_t + (output_t - mean) / rstd * w + b + outputs.append(output_t) + + # Update states for next step + if (t + 1) % chunk_size == 0: + M_prev_nabla = M_t.clone() + M_prev = M_t + S_prev = S_t + + # Stack outputs along sequence dimension + output = torch.stack(outputs, dim=-2).squeeze( + -3, + ) # (batch_size, num_heads, seq_len, dim) + + if output_final_state: + return output, M_prev + return output, None + + +def chunk_titans_linear( + q, k, v, w, b, theta, alpha, eta, eps, chunk_size, initial_state, output_final_state, +): + B, H, T, D = q.shape + num_batch = T // chunk_size + # [num_batch, B, num_heads, mini_batch_size, head_dim] + _q = q.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) + _k = k.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) + _v = v.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) + # [num_batch, B, num_heads, mini_batch_size, 1] + _eta = eta.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) + _theta = theta.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) + _alpha = alpha.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) + # [H, 1, D] + w = w.reshape(H, 1, D).to(torch.float32) + b = b.reshape(H, 1, D).to(torch.float32) + # [num_heads, 1, head_dim] + if initial_state is None: + M_prev = torch.zeros((B, H, D, D), device=v.device, dtype=v.dtype).to( + torch.float32, + ) + else: + M_prev = initial_state + + S_prev = torch.zeros_like(M_prev) + + # [num_batch, B, num_heads, mini_batch_size, head_dim] + o = torch.empty_like(_v) + + for i in range(num_batch): + q_i, k_i, v_i, eta_i, theta_i, alpha_i = [ + x[i] for x in [_q, _k, _v, _eta, _theta, _alpha] + ] + + # beta, beta_T, f, f_T, g, G, m_T, n = combine_params(theta_i, alpha_i, eta_i, chunk_size) + beta, beta_T, f, f_T, g, G, m_T, n = combine_params_log( + theta_i, alpha_i, eta_i, chunk_size, + ) + + m_T = m_T.unsqueeze(-1).unsqueeze(-1) + beta_T = beta_T.unsqueeze(-1).unsqueeze(-1) + f_T = f_T.unsqueeze(-1).unsqueeze(-1) + g_diag = torch.diag_embed(g).to(q_i.dtype) + n = torch.diag_embed(n).to(q_i.dtype) + beta = torch.diag_embed(beta).to(q_i.dtype) + f = torch.diag_embed(f).to(q_i.dtype) + km = k_i @ M_prev + reconstruction_target = v_i - k_i + + mean = km.mean(-1, True) + var = km.var(-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + km_hat = (km - mean) / rstd + + grad = w * km_hat + b - reconstruction_target + grad *= w + v_new = D * grad - grad.sum(-1, keepdim=True) / (rstd * D) + proj_term = km_hat * (grad * km_hat).sum(-1, keepdim=True) / (rstd * D) + v_new = v_new - proj_term + # v_new = (D * grad - grad.sum(-1, True)) + # print(f"Projection term stats: min={torch.abs(beta_T).min()}") + + # v_new = grad + + Attn = torch.tril(q_i @ k_i.transpose(-2, -1)) * G + + # o_i + output_t = beta @ q_i @ M_prev + f @ q_i @ S_prev - 2 * Attn @ v_new + + M_t = ( + beta_T * M_prev + + f_T * S_prev + - 2 * (g_diag @ k_i).transpose(-1, -2) @ v_new + ) + # cal S_T from S_0 + S_t = m_T * S_prev - 2 * (n @ k_i).transpose(-1, -2) @ v_new + # layer norm with residuals + mean = output_t.mean(dim=-1, keepdim=True) + var = output_t.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + output_t = output_t + (output_t - mean) / rstd * w + b + o[i] = output_t + S_prev = S_t + M_prev = M_t + + # [B, num_mini_batch, mini_batch_size, num_heads, head_dim] + o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D) + M_prev = M_prev if output_final_state else None + return o, M_prev + + +# most of the code is copied from ttt +def chunk_titans_linear_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + theta: torch.Tensor, + alpha: torch.Tensor, + eta: torch.Tensor, + eps: float = 1e-6, + chunk_size: int = 16, # chunk size + initial_state: torch.Tensor = None, + output_final_state: bool = False, + head_first: bool = False, + use_chunk: bool = True, +): + assert q.dtype == k.dtype == v.dtype + assert k.shape[-1] == v.shape[-1], "DK must equal to DV." + if not head_first: + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + eta = eta.transpose(1, 2) + alpha = alpha.transpose(1, 2) + theta = theta.transpose(1, 2) + seq_len = q.shape[-2] + pad_len = (chunk_size - (seq_len % chunk_size)) % chunk_size + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + theta = F.pad(theta, (0, 0, 0, pad_len)) + alpha = F.pad(alpha, (0, 0, 0, pad_len)) + eta = F.pad(eta, (0, 0, 0, pad_len)) + theta[:, :, -1, :] = theta[:, :, -(pad_len + 1), :] + alpha[:, :, -1, :] = alpha[:, :, -(pad_len + 1), :] + eta[:, :, -1, :] = eta[:, :, -(pad_len + 1), :] + assert q.shape[-2] % chunk_size == 0, "Sequence length should be a multiple of BT." + q, k, v, w, b = map(lambda x: x.to(torch.float32), [q, k, v, w, b]) + if use_chunk: + o, final_state = chunk_titans_linear( + q, + k, + v, + w, + b, + theta, + alpha, + eta, + eps, + chunk_size, + initial_state, + output_final_state, + ) + else: + o, final_state = titans_linear( + q, + k, + v, + w, + b, + theta, + alpha, + eta, + eps, + chunk_size, + initial_state, + output_final_state, + ) + o = o[:, :, :seq_len, :] + if not head_first: + o = o.transpose(1, 2) + return o, final_state diff --git a/fla/ops/ttt/__init__.py b/fla/ops/ttt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..109b0e425a25bfbaf48802b0c9e74d3b4f62ebbb --- /dev/null +++ b/fla/ops/ttt/__init__.py @@ -0,0 +1,8 @@ + +from .chunk import chunk_ttt_linear +from .fused_chunk import fused_chunk_ttt_linear + +__all__ = [ + 'fused_chunk_ttt_linear', + 'chunk_ttt_linear', +] diff --git a/fla/ops/ttt/chunk.py b/fla/ops/ttt/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..ba62f1c2ce38e26dc622c2c944e5444e711b8842 --- /dev/null +++ b/fla/ops/ttt/chunk.py @@ -0,0 +1,1474 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan + +import warnings + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.modules.layernorm import group_norm +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_ttt_linear_fwd_kernel_h( + k, + v, + v_new, + eta, + w, + b, + eps, + h, + hb, + h0, + hb0, + ht, + hbt, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_hb = tl.zeros([BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + if USE_INITIAL_STATE_B: + p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + + offs = tl.arange(0, BV) + b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) + b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) + + for i_t in range(NT): + p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_hb = tl.make_block_ptr(hb + ((boh + i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_hb, b_hb.to(p_hb.dtype.element_ty), boundary_check=(0,)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + + b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] + b_kh = tl.where((offs < V)[None, :], b_kh, 0.) + mean = tl.sum(b_kh, axis=1, keep_dims=True) / V + xbar = tl.where((offs < V)[None, :], b_kh - mean, 0.) + var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V + rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) + b_kh_hat = (b_kh - mean) * rstd + + b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ + b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) + b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) + b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) + * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + b_eta_last = tl.load(p_eta_last) + b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) + b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_ttt_linear_fwd_kernel_o( + q, + k, + v, + eta, + h, + hb, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + eta += bos * H + i_h + o += (bos * H + i_h) * V + h += (i_tg * H + i_h) * K * V + hb += (i_tg * H + i_h) * V + stride_qk = H*K + stride_vo = H*V + stride_eta = H + + p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (0, i_t * BT), (BK, BT), (0, 1)) + p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) + p_hb = tl.make_block_ptr(hb, (V,), (1,), (i_v * BV,), (BV,), (0,)) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + # [BT, 1] + b_eta = tl.load(p_eta, boundary_check=(0,), padding_option="zero") + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") + # [BV] + b_hb = tl.load(p_hb, boundary_check=(0,), padding_option="zero") + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o = tl.dot(b_q, b_h, allow_tf32=False) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A = tl.dot(b_q, b_k, allow_tf32=False) + + o_i = tl.arange(0, BT) + m_A = o_i[:, None] >= o_i[None, :] + b_A = tl.where(m_A, b_A, 0) + b_Ae = tl.where(m_A, b_eta[:, None], 0.0) + + p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_o = (b_o - tl.dot(b_eta[:, None] * b_A.to(b_v.dtype), b_v, allow_tf32=False)) * scale + b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v.dtype), b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_ttt_linear_bwd_kernel_h( + k, + v, + v_new, + eta, + w, + b, + eps, + h, + h0, + hb0, + x, + y, + r, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NT: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_hb = tl.zeros([BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + if USE_INITIAL_STATE_B: + p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + + offs = tl.arange(0, BV) + b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) + b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) + + for i_t in range(NT): + p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + + b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] + b_kh = tl.where((offs < V)[None, :], b_kh, 0.) + mean = tl.sum(b_kh, axis=1, keep_dims=True) / V + xbar = tl.where((offs < V)[None, :], b_kh - mean, 0.) + var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V + rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) + b_kh_hat = (b_kh - mean) * rstd + + b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ + b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) + b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) + b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) + * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V + tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + b_eta_last = tl.load(p_eta_last) + b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) + b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [4] + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_ttt_linear_bwd_kernel_dv_local( + q, + k, + eta, + do, + dv, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + eta += bos * H + i_h + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + stride_qk = H*K + stride_vo = H*V + stride_eta = H + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, b_q) + + p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) + b_eta = tl.load(p_eta, boundary_check=(0,)) + mask = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]) + b_A = - tl.where(mask, b_A * scale * b_eta[None, :], 0).to(do.dtype.element_ty) + b_Ae = - tl.where(mask, b_eta[None, :], 0).to(do.dtype.element_ty) + + for i_v in range(tl.cdiv(V, BV)): + p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None, + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8, 16] + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_ttt_linear_bwd_kernel_norm( + q, + k, + v, + v_new, + x, + y, + r, + w, + b, + eta, + h, + dht, + dhbt, + dh0, + dhb0, + do, + dh, + dhb, + dv, + dv_new, + dk, + dw, + db, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT_B: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_dhb = tl.zeros([BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") + if USE_FINAL_STATE_GRADIENT_B: + p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") + + # [BV] + offs_v = tl.arange(0, BV) + offs_t = tl.arange(0, BT) + b_w = tl.load(w + i_h * V + offs_v, mask=offs_v < V, other=0.) + b_b = tl.load(b + i_h * V + offs_v, mask=offs_v < V, other=0.) + b_dw = tl.zeros([BV], dtype=b_w.dtype) + b_db = tl.zeros([BV], dtype=b_b.dtype) + p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + + for i_t in range(NT - 1, -1, -1): + p_h = tl.make_block_ptr(h + ((boh+i_t) * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dhb = tl.make_block_ptr(dhb + ((boh+i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dhb, b_dhb.to(p_dhb.dtype.element_ty), boundary_check=(0,)) + p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv_new = tl.make_block_ptr(dv_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, i_k * BK), (BT, BK), (1, 0)) + p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) + p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_dv_new = tl.load(p_dv_new, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_eta_last = tl.load(p_eta_last) + b_dv_new -= tl.dot(b_eta_last * b_k, b_dh.to(b_k.dtype)) + b_dv_new -= b_eta_last * b_dhb.to(b_k.dtype)[None, :] + + b_v_new = tl.load(p_v_new, boundary_check=(0, 1), padding_option="zero") + b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - + b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V + b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + + b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V + b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v_new.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) + + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_w = b_w.to(b_k.dtype) + b_b = b_b.to(b_k.dtype) + b_dv = -b_w * b_dy.to(b_k.dtype) + b_dk = b_w * b_dy.to(b_k.dtype) + b_dw += tl.sum(2 * b_w * b_x * b_dy.to(b_k.dtype) + + (b_b - b_v.to(b_k.dtype) + b_k) * b_dy.to(b_k.dtype), axis=0).to(b_dw.dtype) + b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype) + b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) + + # d_rstd, dx --> dkh --> dk, dh + b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") + b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + b_q = (b_q * scale).to(b_q.dtype) + b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - + b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V + b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V + b_dkh = tl.where((offs_v < V)[None, :] * (offs_t < T-i_t*BT)[:, None], b_dkh, 0.) + b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype) + b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh) + b_dhb += tl.sum(b_do + b_dkh, axis=0) + b_dh = tl.where((offs_v < V)[None, :], b_dh, 0.) + b_dhb = tl.where((offs_v < V), b_dhb, 0.) + + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + if USE_INITIAL_STATE_B: + p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_bwd_kernel_dqke( + q, + k, + v, + e, + h, + do, + dh, + dhb, + dq, + dk, + de, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h += (i_tg * H + i_h) * K * V + dh += (i_tg * H + i_h) * K * V + dhb += (i_tg * H + i_h) * V + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + e += bos * H + i_h + de += bos * H + i_h + stride_qk = H*K + stride_vo = H*V + stride_e = H + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_ds = tl.zeros([BT, BT], dtype=tl.float32) + b_de = tl.zeros([BT], dtype=tl.float32) + + p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + p_e_last = (e + (i_t*BT+BT-1)*stride_e) if (i_t*BT+BT) <= T else (e + (T-1)*stride_e) + i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1) + mask = (tl.arange(0, BT) == i_last) + b_e_last = tl.load(p_e_last) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dhb = tl.make_block_ptr(dhb, (V,), (1,), (i_v * BV,), (BV,), (0,)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BV] + b_dhb = tl.load(p_dhb, boundary_check=(0,)) + # [BT, BV] @ [BV, BT] -> [BT, BT] + b_ds += tl.dot(b_do, tl.trans(b_v)) + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + # [BT, BV] @ [BV, BK] -> [BT, BK] + b_dk -= b_e_last * tl.dot(b_v, b_dh.to(b_v.dtype)) + b_de -= mask * tl.sum(tl.trans(b_dh) * tl.dot(tl.trans(b_k), b_v.to(b_k.dtype))) + b_de -= mask * tl.sum(b_dhb * tl.sum(b_v, axis=0).to(b_k.dtype)) + + o_i = tl.arange(0, BT) + p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_e = tl.make_block_ptr(e, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_e = tl.load(p_e, boundary_check=(0,)) + + p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_de = tl.make_block_ptr(de, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) + + b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds, 0) + b_ds = b_ds.to(b_k.dtype) + b_dq -= tl.dot(b_ds, b_k) * b_e[:, None] + b_dk -= tl.dot(tl.trans(b_ds), b_q * b_e[:, None]) * scale + b_de -= tl.sum(scale * tl.dot(b_ds, b_k) * b_q, axis=1) + b_de -= tl.sum(b_ds, axis=1) + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) + + +def chunk_ttt_linear_fwd_h( + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + eps: float, + initial_state: torch.Tensor | None = None, + initial_state_bias: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + assert NV == 1, 'NV > 1 is not supported by TTT update rule.' + + h = k.new_empty(B, NT, H, K, V) + hb = k.new_empty(B, NT, H, 1, V) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(v) + grid = (NK, NV, N * H) + + chunk_ttt_linear_fwd_kernel_h[grid]( + k=k, + v=v, + v_new=v_new, + eta=eta, + w=w, + b=b, + eps=eps, + h=h, + hb=hb, + h0=initial_state, + hb0=initial_state_bias, + ht=final_state, + hbt=final_state_bias, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return h, hb, v_new, final_state, final_state_bias + + +def chunk_ttt_linear_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + eta: torch.Tensor, + h: torch.Tensor, + hb: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K, V = *q.shape, v.shape[-1] + if scale is None: + scale = k.shape[-1] ** -0.5 + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + assert NV == 1, 'NV > 1 is not supported by TTT update rule.' + + o = torch.empty_like(v) + + grid = (NV, NT, B * H) + chunk_ttt_linear_fwd_kernel_o[grid]( + q, + k, + v, + eta, + h, + hb, + o, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return o + + +def chunk_ttt_linear_bwd_h( + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + eps: float, + initial_state: torch.Tensor | None = None, + initial_state_bias: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' + assert NV == 1, 'NV > 1 is not supported by TTT update rule.' + + h = k.new_empty(B, NT, H, K, V) + rstd = v.new_empty(B, T, H, 1, dtype=torch.float32) + x = torch.empty_like(v) + y = torch.empty_like(v) + + v_new = torch.empty_like(v) + grid = (NK, NV, N * H) + + chunk_ttt_linear_bwd_kernel_h[grid]( + k=k, + v=v, + v_new=v_new, + eta=eta, + w=w, + b=b, + eps=eps, + h=h, + h0=initial_state, + hb0=initial_state_bias, + x=x, + y=y, + r=rstd, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + NT=NT, + ) + return h, v_new, x, y, rstd + + +def chunk_ttt_linear_bwd_dv_local( + q: torch.Tensor, + k: torch.Tensor, + eta: torch.Tensor, + do: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = min(max(triton.next_power_of_2(K), 16), 128) + BV = min(max(triton.next_power_of_2(V), 16), 128) + + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_ttt_linear_bwd_kernel_dv_local[grid]( + q, + k, + eta, + do, + dv, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dv + + +def chunk_ttt_linear_bwd_norm( + q: torch.Tensor, # [B, H, L, D] + k: torch.Tensor, # [B, H, L, D] + v: torch.Tensor, # [B, H, L, D] + v_new: torch.Tensor, # [B, H, L, D] + x: torch.Tensor, # [B, H, L, D] + y: torch.Tensor, # [B, H, L, D] + rstd: torch.Tensor, # [B, H, L, 1] + w: torch.Tensor, # [H, D] + b: torch.Tensor, # [H, D] + eta: torch.Tensor, # [B, H, L, 1] + h0: torch.Tensor, # [B, H, D, D] + hb0: torch.Tensor, # [B, H, 1, D] + h: torch.Tensor, # [B, H, NT, D, D] + dht: torch.Tensor | None, # [B, H, D, D] + dhbt: torch.Tensor | None, # [B, H, 1, D] + dv_new: torch.Tensor | None, # [B, H, L, D] + do: torch.Tensor, # [B, H, L, D] + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # torch implementation of `dkh, dw, db, dk, dv` for LN^2 + assert cu_seqlens is None, "bwd of varlen is not implemented yet." + B, T, H, K, V = *q.shape, do.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + + BK = max(triton.next_power_of_2(K), 16) + BV = max(triton.next_power_of_2(V), 16) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, 'NK > 1 is not supported by TTT.' + assert NV == 1, 'NV > 1 is not supported by TTT.' + + dh = q.new_empty(B, NT, H, K, V) + dhb = q.new_empty(B, NT, H, 1, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dhb0 = torch.empty_like(hb0, dtype=torch.float32) if hb0 is not None else None + dv = torch.empty_like(v) + dk = torch.empty_like(k) + dw = w.new_empty(B, H, V) + db = b.new_empty(B, H, V) + + grid = (NK, NV, N * H) + chunk_ttt_linear_bwd_kernel_norm[grid]( + q=q, + k=k, + v=v, + v_new=v_new, + x=x, + y=y, + r=rstd, + w=w, + b=b, + eta=eta, + h=h, + dht=dht, + dhbt=dhbt, + dh0=dh0, + dhb0=dhb0, + do=do, + dh=dh, + dhb=dhb, + dv=dv, + dv_new=dv_new, + dk=dk, + dw=dw, + db=db, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dw = dw.sum(dim=0) + db = db.sum(dim=0) + return dh, dhb, dh0, dhb0, dv, dk, dw, db + + +def chunk_ttt_linear_bwd_norm_ref( + q: torch.Tensor, # [B, H, L, D] + k: torch.Tensor, # [B, H, L, D] + v: torch.Tensor, # [B, H, L, D] + v_new: torch.Tensor, # [B, H, L, D] + kh: torch.Tensor, # [B, H, L, D] + y: torch.Tensor, # [B, H, L, D] + w: torch.Tensor, # [H, D] + b: torch.Tensor, # [H, D] + eta: torch.Tensor, # [B, H, L, 1] + h0: torch.Tensor, # [B, H, D, D] + h: torch.Tensor, # [B, H, NT, D, D] + dht: torch.Tensor | None, # [B, H, D, D] + dv_new: torch.Tensor | None, # [B, H, L, D] + do: torch.Tensor, # [B, H, L, D] + scale: float, + eps: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # torch implementation of `dkh, dw, db, dk, dv` for LN^2 + assert cu_seqlens is None, "bwd of varlen is not implemented yet." + B, T, H, K, V = *q.shape, do.shape[-1] + # [B, L, H, D] -> [B, H, L, D] + q, k, v, v_new, kh, y, h, eta, dv_new, do = [ + x.transpose(1, 2) for x in + [q, k, v, v_new, kh, y, h, eta, dv_new, do] + ] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q, k, v, v_new, kh, y, eta, dv_new, do = [ + F.pad(x, (0, 0, 0, pad_len)) for x in + [q, k, v, v_new, kh, y, eta, dv_new, do] + ] + eta[:, :, -1, :] = eta[:, :, -(pad_len+1), :] + # [NT, B, H, BT, D] + q, k, v, v_new, kh, y, eta, dv_new, do = [ + x.reshape(B, H, NT, BT, -1).permute(2, 0, 1, 3, 4) for x in + [q, k, v, v_new, kh, y, eta, dv_new, do] + ] + h = h.permute(2, 0, 1, 3, 4) + + # allocate + dh = q.new_zeros(NT, B, H, K, V) + dv = torch.zeros_like(v) + dk = torch.zeros_like(k) + dw = torch.zeros_like(w) + db = torch.zeros_like(b) + # recurrent state + b_dh = dht if dht is not None else torch.zeros_like(dh[0]) + b_dh = b_dh.to(torch.float32) + + # [H, 1, D] + _w = w.reshape(H, 1, V).to(torch.float32) + _b = b.reshape(H, 1, V).to(torch.float32) + + # d_state passing + for i_t in range(NT - 1, -1, -1): + dh[i_t] = b_dh.to(dh.dtype) + # [B, H, BT, D] + _q, _k, _v, _v_new, _kh, _y, _h, _eta, _dv_new, _do = [ + x[i_t].to(torch.float32) for x in + (q, k, v, v_new, kh, y, h, eta, dv_new, do) + ] + _dv_new -= (_eta[:, :, -1, :, None] * _k) @ b_dh + + mean = _kh.mean(dim=-1, keepdim=True) + var = _kh.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = 1 / torch.sqrt(var + eps).to(torch.float32) + x = (_kh - mean) * rstd + # [B, H, BT, D] + dy = rstd * (_dv_new*V - _dv_new.sum(dim=-1, keepdim=True) - x*(x*_dv_new).sum(dim=-1, keepdim=True)) / V + dx = -rstd * (_dv_new*(x*_y).sum(dim=-1, keepdim=True) + _y*(x*_dv_new).sum(dim=-1, keepdim=True)) / V + d_rstd = (_dv_new * _v_new / rstd).sum(dim=-1, keepdim=True) + + dv[i_t] = (-_w*dy).to(dv.dtype) + dk[i_t] += (_w*dy).to(dk.dtype) + dw += (2*_w*x*dy+(_b-_v+_k)*dy).sum(dim=(0, 2)).to(dw.dtype) + db += (_w*dy).sum(dim=(0, 2)).to(db.dtype) + dx += _w*_w*dy + + # d_rstd, dx --> dkh --> dk, dh + dkh = rstd * (V * dx - dx.sum(dim=-1, keepdim=True) - x * (x * dx).sum(dim=-1, keepdim=True)) / V + dkh -= rstd**2 * d_rstd * x / V + dk[i_t] += (dkh @ _h.transpose(-2, -1)).to(dk.dtype) + b_dh += (_q.transpose(-2, -1) * scale) @ _do + _k.transpose(-2, -1) @ dkh + dh0 = b_dh.to(torch.float32) if h0 is not None else None + + # [NT, B, H, BT, D] -> [B, H, T, D] + dv = dv.permute(1, 2, 0, 3, 4).reshape(B, H, -1, V)[:, :, :T, :] + dk = dk.permute(1, 2, 0, 3, 4).reshape(B, H, -1, K)[:, :, :T, :] + # [B, H, NT, D, D] + dh = dh.permute(1, 2, 0, 3, 4) + dv, dk, dh = [x.transpose(1, 2) for x in (dv, dk, dh)] + dh, dv, dk, dw, db = [x.contiguous() for x in (dh, dv, dk, dw, db)] + dh0 = dh0.contiguous() if h0 is not None else None + return dh, dh0, dv, dk, dw, db + + +def chunk_ttt_linear_bwd_dqke( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + eta: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dhb: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BK = max(triton.next_power_of_2(K), 16) + BV = min(max(triton.next_power_of_2(V), 16), 64) + NK = triton.cdiv(K, BK) + assert NK == 1, "NK > 1 is not supported." + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + de = torch.empty_like(eta) + grid = (NK, NT, B * H) + + chunk_bwd_kernel_dqke[grid]( + q=q, + k=k, + v=v, + e=eta, + h=h, + do=do, + dh=dh, + dhb=dhb, + dq=dq, + dk=dk, + de=de, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dq, dk, de + + +def chunk_ttt_linear_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + initial_state: torch.Tensor, + initial_state_bias: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 16, + chunk_indices: torch.LongTensor | None = None, +): + BT = chunk_size + h, hb, v_new, final_state, final_state_bias = chunk_ttt_linear_fwd_h( + k=k, + v=v, + w=w, + b=b, + eta=eta, + eps=eps, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + o = chunk_ttt_linear_fwd_o( + q=q, + k=k, + v=v_new, + eta=eta, + h=h, + hb=hb, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + return o, final_state, final_state_bias + + +def chunk_ttt_linear_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + do: torch.Tensor, + dht: torch.Tensor, + dhbt: torch.Tensor, + chunk_size: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +): + BT = chunk_size + h, v_new, x, y, rstd = chunk_ttt_linear_bwd_h( + k=k, + v=v, + w=w, + b=b, + eta=eta, + eps=eps, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + dv_new = chunk_ttt_linear_bwd_dv_local( + q=q, + k=k, + eta=eta, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + dh, dhb, dh0, dhb0, dv, dk, dw, db = chunk_ttt_linear_bwd_norm( + q=q, + k=k, + v=v, + v_new=v_new, + x=x, + y=y, + rstd=rstd, + w=w, + b=b, + eta=eta, + h0=initial_state, + hb0=initial_state_bias, + h=h, + dht=dht, + dhbt=dhbt, + dv_new=dv_new, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + dq, dk2, de = chunk_ttt_linear_bwd_dqke( + q=q, + k=k, + v=v_new, + eta=eta, + h=h, + do=do, + dh=dh, + dhb=dhb, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + return dq, dk, dv, de, dw, db, dh0, dhb0 + + +class ChunkTTTLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + w, + b, + chunk_size, + eta, + scale, + eps, + initial_state, + initial_state_bias, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ): + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + o, final_state, final_state_bias = chunk_ttt_linear_fwd( + q=q, + k=k, + v=v, + w=w, + b=b, + eta=eta, + scale=scale, + eps=eps, + chunk_size=chunk_size, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias, chunk_indices) + ctx.chunk_size = chunk_size + ctx.scale = scale + ctx.eps = eps + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), final_state, final_state_bias + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht, dhbt): + q, k, v, eta, w, b, initial_state, initial_state_bias, chunk_indices = ctx.saved_tensors + dq, dk, dv, de, dw, db, dh0, dhb0 = chunk_ttt_linear_bwd( + q=q, + k=k, + v=v, + w=w, + b=b, + eta=eta, + scale=ctx.scale, + eps=ctx.eps, + do=do, + dht=dht, + dhbt=dhbt, + chunk_size=ctx.chunk_size, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + cu_seqlens=ctx.cu_seqlens, + chunk_indices=chunk_indices, + ) + return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None, None, None + + +def norm_residual(x, weight, bias, eps): + # GroupNorm and Residual + B, T, H, D = x.shape + x += group_norm( + x.reshape(B, T, -1).clone(), + weight=weight.reshape(-1).clone(), + bias=bias.reshape(-1).clone(), + eps=eps, + num_groups=H, + ).reshape(x.shape) + return x + + +def chunk_ttt_linear( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float = None, + eps: float = 1e-6, + chunk_size: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + head_first: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `(B, H, T, K)` + k (torch.Tensor): + keys of shape `(B, H, T, K)` + v (torch.Tensor): + values of shape `(B, H, T, V)` + w (torch.Tensor): + layer norm weight of shape `(H, V)` + b (torch.Tensor): + layer norm bias of shape `(H, V)` + eta (torch.Tensor): + Learning rate for hidden state, of shape `(B, H, T, 1)`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + chunk_size (int): + chunk size. Default: `16`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `(B, H, K, V)`. Default: `None`. + initial_state_bias (Optional[torch.Tensor]): + Initial state bias of shape `(B, H, 1, V)`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `(B, H, K, V)`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` + final_state (torch.Tensor): + Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None` + """ + assert q.dtype == k.dtype == v.dtype + assert k.shape[-1] == v.shape[-1], "DK must equal to DV." + if isinstance(eta, float): + eta = torch.full_like(q[:, :, :, :1], eta) + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "Scale must be positive." + o, final_state, final_state_bias = ChunkTTTLinearFunction.apply( + q, + k, + v, + w, + b, + chunk_size, + eta, + scale, + eps, + initial_state, + initial_state_bias, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + ) + o = norm_residual(o, w, b, eps) + return o, final_state, final_state_bias diff --git a/fla/ops/ttt/fused_chunk.py b/fla/ops/ttt/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..ca300ff026c010310a5b9a64fe6a0981993aa4e3 --- /dev/null +++ b/fla/ops/ttt/fused_chunk.py @@ -0,0 +1,832 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.modules.layernorm import group_norm +from fla.utils import IS_NVIDIA_HOPPER, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + +NUM_WARPS = [1, 2] if IS_NVIDIA_HOPPER else [1, 2, 4, 8] + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_chunk_ttt_linear_fwd_kernel( + q, + k, + v, + eta, + w, + b, + o, + scale, + eps, + h0, + hb0, + ht, + hbt, + cu_seqlens, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + + o_i = tl.arange(0, BT) + v_i = tl.arange(0, BV) + m_A = o_i[:, None] >= o_i[None, :] + b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) + b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_hb = tl.zeros([BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + if USE_INITIAL_STATE_B: + p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + + for i_t in range(NT): + p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + + # [BT, BV] + b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] + b_kh = tl.where((v_i < V)[None, :], b_kh, 0.) + mean = tl.sum(b_kh, axis=1, keep_dims=True) / V + xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.) + var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V + rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) + b_kh_hat = (b_kh - mean) * rstd + + b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ + b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) + b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) + b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) + * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + # [BT] + b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") + b_q = (b_q * scale).to(b_k.dtype) + + # [BT, BT] + b_A = tl.dot(b_q, b_k, allow_tf32=False) + b_A = tl.where(m_A, b_A, 0) + b_Ae = tl.where(m_A, b_e[:, None], 0.0) + + b_o = - tl.dot(b_e[:, None] * b_A.to(b_v2.dtype), b_v2, allow_tf32=False) + b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v2.dtype), b_v2, allow_tf32=False) + b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False) + b_e_last = tl.load(p_e_last) + b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) + b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) + b_h = tl.where((v_i < V)[None, :], b_h, 0.) + b_hb = tl.where((v_i < V), b_hb, 0.) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_chunk_ttt_linear_bwd_kernel_h( + k, + v, + v2, + x, + y, + r, + w, + b, + eta, + h0, + hb0, + h, + do, + dq, + scale, + eps, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, +): + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + bos, _ = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + o_i = tl.arange(0, BT) + v_i = tl.arange(0, BV) + m_A = o_i[:, None] >= o_i[None, :] + b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) + b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) + + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_hb = tl.zeros([BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + if USE_INITIAL_STATE_B: + p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + + for i_t in range(NT): + p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) + p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + p_dq = tl.make_block_ptr(dq+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) + p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + + b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] + b_kh = tl.where((v_i < V)[None, :], b_kh, 0.) + mean = tl.sum(b_kh, axis=1, keep_dims=True) / V + xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.) + var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V + rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) + b_kh_hat = (b_kh - mean) * rstd + + b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ + b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) + b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) + b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) + * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V + tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_v2, b_v2.to(p_v2.dtype.element_ty), boundary_check=(0, 1)) + + b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") + b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + + b_v2 = tl.where((v_i < V)[None, :], b_v2, 0.) + b_ds = tl.dot(b_do, tl.trans(b_v2).to(b_do.dtype)) + b_ds = tl.where(m_A, b_ds, 0) + b_ds = b_ds.to(b_k.dtype) + b_dq = tl.dot(b_do, tl.trans(b_h).to(b_do.dtype)) + b_dq -= tl.dot(b_ds, tl.trans(b_k)) * b_e[:, None] + b_dq *= scale + + b_e_last = tl.load(p_e_last) + b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) + b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) + b_h = tl.where((v_i < V)[None, :], b_h, 0.) + b_hb = tl.where((v_i < V), b_hb, 0.) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, + 'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS + ], + key=['BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_chunk_ttt_linear_bwd_kernel_dh( + q, + k, + v, + v2, + x, + y, + r, + w, + b, + eta, + h, + dht, + dhbt, + dh0, + dhb0, + do, + dk, + dv, + de, + dw, + db, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_STATE_B: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + USE_FINAL_STATE_GRADIENT_B: tl.constexpr, +): + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + bos, _ = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + # [BV] + b_dhb = tl.zeros([BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") + if USE_FINAL_STATE_GRADIENT_B: + p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") + + # [BV] + o_i = tl.arange(0, BT) + v_i = tl.arange(0, BV) + m_A = o_i[:, None] >= o_i[None, :] + m_A_t = o_i[:, None] <= o_i[None, :] + b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) + b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) + b_dw = tl.zeros([BV], dtype=b_w.dtype) + b_db = tl.zeros([BV], dtype=b_b.dtype) + p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + + for i_t in range(NT - 1, -1, -1): + p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) + p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) + p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) + p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + p_de = tl.make_block_ptr(de+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H + b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") + b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + b_e_last = tl.load(p_e_last) + b_A = tl.dot(b_k, b_q) + b_A = - tl.where(m_A_t, b_A * scale * b_e[None, :], 0).to(do.dtype.element_ty) + b_Ae = - tl.where(m_A_t, b_e[None, :], 0).to(do.dtype.element_ty) + b_dv_new = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do) + b_dv_new -= tl.dot(b_e_last * b_k, b_dh.to(b_k.dtype)) + b_dv_new -= b_e_last * b_dhb.to(b_k.dtype)[None, :] + + b_v2 = tl.load(p_v2, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - + b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V + b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + + b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V + b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v2.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) + + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_w = b_w.to(b_k.dtype) + b_b = b_b.to(b_k.dtype) + b_dv = -b_w * b_dy.to(b_k.dtype) + b_dk = b_w * b_dy.to(b_k.dtype) + b_dw += tl.sum(2 * b_w * b_x * b_dy.to(b_k.dtype) + + (b_b - b_v.to(b_k.dtype) + b_k) * b_dy.to(b_k.dtype), axis=0).to(b_dw.dtype) + b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype) + b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) + + b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") + b_q = (b_q * scale).to(b_q.dtype) + b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - + b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V + b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V + b_dkh = tl.where((v_i < V)[None, :] * (o_i < T-i_t*BT)[:, None], b_dkh, 0.) + b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype) + + b_ds = tl.dot(b_do, tl.trans(b_v2)) + b_ds = tl.where(m_A, b_ds, 0) + b_ds = b_ds.to(b_k.dtype) + i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1) + mask = (o_i == i_last) + b_dk -= b_e_last * tl.dot(b_v2, tl.trans(b_dh).to(b_v2.dtype)) + b_dk -= tl.dot(tl.trans(b_ds), tl.trans(b_q) * b_e[:, None]) + b_de = mask * tl.sum(- b_dh * tl.trans(tl.dot(tl.trans(b_v2), b_k))).to(b_k.dtype) + b_de -= mask * tl.sum(b_dhb * tl.sum(b_v2, axis=0)).to(b_k.dtype) + b_de -= tl.sum(tl.dot(b_ds, b_k) * tl.trans(b_q).to(b_k.dtype), axis=1) + b_de -= tl.sum(b_ds, axis=1) + b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh) + b_dhb += tl.sum(b_do + b_dkh, axis=0) + b_dh = tl.where((v_i < V)[None, :], b_dh, 0.) + b_dhb = tl.where((v_i < V), b_dhb, 0.) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + if USE_INITIAL_STATE: + p_dh0 = tl.make_block_ptr(dh0+i_nh*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + if USE_INITIAL_STATE_B: + p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (0,), (BV,), (0,)) + tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) + + +def fused_chunk_ttt_linear_bwd_h( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + do: torch.Tensor, + BT: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, +): + assert cu_seqlens is None, "bwd of varlen is not implemented yet." + B, T, H, K, V = *k.shape, v.shape[-1] + # N: the actual number of sequences in the batch with either equal or variable lengths + N, NT = B, triton.cdiv(T, BT) + BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) + assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." + + h = k.new_empty(B, NT, H, K, V) + r = v.new_empty(B, T, H, 1, dtype=torch.float32) + v2 = torch.empty_like(v) + x = torch.empty_like(v) + y = torch.empty_like(v) + dq = torch.empty_like(q) + + grid = (N * H,) + fused_chunk_ttt_linear_bwd_kernel_h[grid]( + k=k, + v=v, + v2=v2, + x=x, + y=y, + r=r, + w=w, + b=b, + eta=eta, + h0=initial_state, + hb0=initial_state_bias, + h=h, + do=do, + dq=dq, + scale=scale, + eps=eps, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dq, h, v2, x, y, r + + +def fused_chunk_ttt_linear_bwd_dh( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + v2: torch.Tensor, + x: torch.Tensor, + y: torch.Tensor, + r: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + h: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dhbt: torch.Tensor, + BT: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, +): + assert cu_seqlens is None, "bwd of varlen is not implemented yet." + B, T, H, K, V = *k.shape, v.shape[-1] + # N: the actual number of sequences in the batch with either equal or variable lengths + N = B + BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) + assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." + + dh0 = torch.empty_like(initial_state, dtype=torch.float32) if initial_state is not None else None + dhb0 = torch.empty_like(initial_state_bias, dtype=torch.float32) if initial_state_bias is not None else None + dk = torch.empty_like(k) + dv = torch.empty_like(v) + de = torch.empty_like(eta) + dw = w.new_empty(B, H, V) + db = b.new_empty(B, H, V) + + grid = (N * H,) + fused_chunk_ttt_linear_bwd_kernel_dh[grid]( + q=q, + k=k, + v=v, + v2=v2, + x=x, + y=y, + r=r, + w=w, + b=b, + eta=eta, + h=h, + dht=dht, + dhbt=dhbt, + dh0=dh0, + dhb0=dhb0, + do=do, + dk=dk, + dv=dv, + de=de, + dw=dw, + db=db, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + dw = dw.sum(dim=0) + db = db.sum(dim=0) + return dk, dv, de, dw, db, dh0, dhb0 + + +def fused_chunk_ttt_linear_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + initial_state: torch.Tensor, + initial_state_bias: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + BT: int = 16, +): + B, T, H, K, V = *k.shape, v.shape[-1] + # N: the actual number of sequences in the batch with either equal or variable lengths + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) + assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." + o = torch.empty_like(v) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None + + grid = (N * H,) + fused_chunk_ttt_linear_fwd_kernel[grid]( + q=q, + k=k, + v=v, + eta=eta, + w=w, + b=b, + o=o, + scale=scale, + eps=eps, + h0=initial_state, + hb0=initial_state_bias, + ht=final_state, + hbt=final_state_bias, + cu_seqlens=cu_seqlens, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return o, final_state, final_state_bias + + +def fused_chunk_ttt_linear_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + do: torch.Tensor, + dht: torch.Tensor, + dhbt: torch.Tensor, + BT: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, +): + assert cu_seqlens is None, "bwd of varlen is not implemented yet." + dq, h, v2, x, y, rstd = fused_chunk_ttt_linear_bwd_h( + q=q, + k=k, + v=v, + w=w, + b=b, + eta=eta, + scale=scale, + eps=eps, + do=do, + BT=BT, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + cu_seqlens=cu_seqlens, + ) + dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd_dh( + q=q, + k=k, + v=v, + v2=v2, + x=x, + y=y, + r=rstd, + w=w, + b=b, + eta=eta, + scale=scale, + h=h, + do=do, + dht=dht, + dhbt=dhbt, + BT=BT, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + cu_seqlens=cu_seqlens, + ) + return dq, dk, dv, de, dw, db, dh0, dhb0 + + +class FusedChunkTTTLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state, + initial_state_bias, output_final_state, cu_seqlens): + o, final_state, final_state_bias = fused_chunk_ttt_linear_fwd( + q=q, + k=k, + v=v, + w=w, + b=b, + eta=eta, + scale=scale, + eps=eps, + BT=BT, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias) + ctx.BT = BT + ctx.scale = scale + ctx.eps = eps + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), final_state, final_state_bias + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht, dhbt): + q, k, v, eta, w, b, initial_state, initial_state_bias = ctx.saved_tensors + dq, dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd( + q=q, + k=k, + v=v, + w=w, + b=b, + eta=eta, + scale=ctx.scale, + eps=ctx.eps, + do=do, + dht=dht, + dhbt=dhbt, + BT=ctx.BT, + initial_state=initial_state, + initial_state_bias=initial_state_bias, + cu_seqlens=ctx.cu_seqlens, + ) + return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None + + +def norm_residual(x, weight, bias, eps): + # GroupNorm and Residual + B, T, H, D = x.shape + x += group_norm( + x.reshape(B, T, -1).clone(), + weight=weight.reshape(-1).clone(), + bias=bias.reshape(-1).clone(), + eps=eps, + num_groups=H, + ).reshape(x.shape) + return x + + +def fused_chunk_ttt_linear( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float = None, + eps: float = 1e-6, + chunk_size: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `(B, H, T, K)` + k (torch.Tensor): + keys of shape `(B, H, T, K)` + v (torch.Tensor): + values of shape `(B, H, T, V)` + w (torch.Tensor): + layer norm weight of shape `(H, V)` + b (torch.Tensor): + layer norm bias of shape `(H, V)` + eta (torch.Tensor): + Learning rate for hidden state, of shape `(B, H, T, 1)`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + chunk_size (int): + chunk size. Default: `16`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `(B, H, K, V)`. Default: `None`. + initial_state_bias (Optional[torch.Tensor]): + Initial state bias of shape `(B, H, 1, V)`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `(B, H, K, V)`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + This argument has been deprecated. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` + final_state (torch.Tensor): + Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None`. + final_state_bias (torch.Tensor): + Final state bias of shape `[B, H, 1, V]` if `output_final_state=True` else `None`. + """ + assert q.dtype == k.dtype == v.dtype + assert k.shape[-1] == v.shape[-1], "DK must equal to DV." + if isinstance(eta, float): + eta = torch.full_like(q[:, :, :, :1], eta) + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead.", + ) + if not head_first and q.shape[1] < q.shape[2]: + warnings.warn( + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + "when head_first=False was specified. " + "Please verify your input tensor format matches the expected shape [B, T, H, ...].", + ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "Scale must be positive." + o, final_state, final_state_bias = FusedChunkTTTLinearFunction.apply( + q, + k, + v, + w, + b, + chunk_size, + eta, + scale, + eps, + initial_state, + initial_state_bias, + output_final_state, + cu_seqlens, + ) + o = norm_residual(o, w, b, eps) + return o, final_state, final_state_bias diff --git a/fla/ops/ttt/naive.py b/fla/ops/ttt/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..7681f11d8d553a18594930a6865bbf1a87e54325 --- /dev/null +++ b/fla/ops/ttt/naive.py @@ -0,0 +1,125 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan + +import torch +import torch.nn.functional as F + + +def ttt_linear( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float, + eps: float, + mini_batch_size: int, + initial_state: torch.Tensor, + initial_state_bias: torch.Tensor, + output_final_state: bool, +): + B, H, T, D = q.shape + BT = mini_batch_size + NT = T // BT + # [NT, B, H, mini_batch_size, D] + _q = q.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) + _k = k.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) + _v = v.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) + # [NT, B, H, BT, 1] + _eta = eta.reshape(B, H, NT, BT, 1).permute(2, 0, 1, 3, 4) + # [H, 1, D] + w = w.reshape(H, 1, D).to(torch.float32) + b = b.reshape(H, 1, D).to(torch.float32) + + h = torch.zeros((B, H, D, D), device=v.device, dtype=torch.float32) if initial_state is None else initial_state + hb = torch.zeros((B, H, 1, D), device=v.device, dtype=torch.float32) if initial_state_bias is None else initial_state_bias + q *= scale + # [NT, B, H, BT, D] + o = torch.empty_like(_v) + + for i in range(NT): + q_i, k_i, v_i, eta_i = [x[i] for x in [_q, _k, _v, _eta]] + kh = k_i @ h + hb + reconstruction_target = v_i - k_i + + mean = kh.mean(-1, True) + var = kh.var(-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + kh_hat = (kh - mean) / rstd + + g = w * kh_hat + b - reconstruction_target + g *= w + v_new = (D * g - g.sum(-1, True) - kh_hat * (g * kh_hat).sum(-1, True)) / (rstd * D) + + Attn = torch.tril(q_i @ k_i.transpose(-2, -1)) + o_i = q_i @ h - (eta_i * Attn) @ v_new + hb - torch.tril(eta_i.expand_as(Attn)) @ v_new + h = h - (eta_i[:, :, -1, :, None] * k_i).transpose(-1, -2) @ v_new + hb = hb - torch.sum(eta_i[:, :, -1, :, None] * v_new, dim=-2, keepdim=True) + # layer norm with residuals + + mean = o_i.mean(dim=-1, keepdim=True) + var = o_i.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) + rstd = torch.sqrt(var + eps).to(torch.float32) + o[i] = o_i + (o_i - mean) / rstd * w + b + + # [B, H, T, D] + o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D) + h = h if output_final_state else None + hb = hb if output_final_state else None + return o, h, hb + + +def chunk_ttt_linear_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eta: torch.Tensor, + scale: float = None, + eps: float = 1e-6, + mini_batch_size: int = 16, + initial_state: torch.Tensor = None, + initial_state_bias: torch.Tensor = None, + output_final_state: bool = False, + head_first: bool = False, +): + assert q.dtype == k.dtype == v.dtype + assert k.shape[-1] == v.shape[-1], "The key and value dimension must be the same." + if isinstance(eta, float): + eta = torch.full_like(q[:, :, :, :1], eta) + if scale is None: + scale = k.shape[-1] ** -0.5 + if not head_first: + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + eta = eta.transpose(1, 2) + T = q.shape[-2] + padded = (mini_batch_size - (T % mini_batch_size)) % mini_batch_size + if padded > 0: + q = F.pad(q, (0, 0, 0, padded)) + k = F.pad(k, (0, 0, 0, padded)) + v = F.pad(v, (0, 0, 0, padded)) + eta = F.pad(eta, (0, 0, 0, padded)) + eta[:, :, -1, :] = eta[:, :, -(padded+1), :] + assert q.shape[-2] % mini_batch_size == 0, "Sequence length should be a multiple of mini_batch_size." + q, k, v, eta, w, b = map(lambda x: x.to(torch.float32), [q, k, v, eta, w, b]) + o, final_state, final_state_bias = ttt_linear( + q, + k, + v, + w, + b, + eta, + scale, + eps, + mini_batch_size, + initial_state, + initial_state_bias, + output_final_state, + ) + o = o[:, :, :T, :].contiguous() + if not head_first: + o = o.transpose(1, 2) + return o, final_state, final_state_bias diff --git a/fla/ops/utils/__init__.py b/fla/ops/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..380a2f08d02594da67f6237585f8412540363796 --- /dev/null +++ b/fla/ops/utils/__init__.py @@ -0,0 +1,56 @@ +from .cumsum import ( + chunk_global_cumsum, + chunk_global_cumsum_scalar, + chunk_global_cumsum_vector, + chunk_local_cumsum, + chunk_local_cumsum_scalar, + chunk_local_cumsum_vector, +) +from .index import ( + get_max_num_splits, + prepare_chunk_indices, + prepare_chunk_offsets, + prepare_cu_seqlens_from_lens, + prepare_cu_seqlens_from_mask, + prepare_lens, + prepare_lens_from_mask, + prepare_position_ids, + prepare_sequence_ids, + prepare_token_indices, +) +from .logsumexp import logsumexp_fwd +from .matmul import addmm, matmul +from .pack import pack_sequence, unpack_sequence +from .pooling import mean_pooling +from .softmax import softmax_bwd, softmax_fwd +from .softplus import softplus +from .solve_tril import solve_tril + +__all__ = [ + "addmm", + "chunk_global_cumsum", + "chunk_global_cumsum_scalar", + "chunk_global_cumsum_vector", + "chunk_local_cumsum", + "chunk_local_cumsum_scalar", + "chunk_local_cumsum_vector", + "get_max_num_splits", + "logsumexp_fwd", + "matmul", + "mean_pooling", + "pack_sequence", + "prepare_chunk_indices", + "prepare_chunk_offsets", + "prepare_cu_seqlens_from_lens", + "prepare_cu_seqlens_from_mask", + "prepare_lens", + "prepare_lens_from_mask", + "prepare_position_ids", + "prepare_sequence_ids", + "prepare_token_indices", + "softmax_bwd", + "softmax_fwd", + "softplus", + "solve_tril", + "unpack_sequence", +] diff --git a/fla/ops/utils/constant.py b/fla/ops/utils/constant.py new file mode 100644 index 0000000000000000000000000000000000000000..ce7a2ccb15ea34f9b7002f74f71a35922bfa08f8 --- /dev/null +++ b/fla/ops/utils/constant.py @@ -0,0 +1,5 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Approximate value of 1/ln(2), used for log/exp base conversion +# Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B) +RCP_LN2 = 1.4426950216 diff --git a/fla/ops/utils/cumsum.py b/fla/ops/utils/cumsum.py new file mode 100644 index 0000000000000000000000000000000000000000..906f37b3be83bea6d9a71e57ddf0de8d2f5e521e --- /dev/null +++ b/fla/ops/utils/cumsum.py @@ -0,0 +1,469 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autotune_cache_kwargs, check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + else: + p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BS': BS}, num_warps=num_warps) + for BS in BS_LIST + for num_warps in [2, 4, 8] + ], + key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_local_cumsum_vector_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + else: + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if REVERSE: + b_o = tl.cumsum(b_s, axis=0, reverse=True) + else: + b_o = tl.cumsum(b_s, axis=0) + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for BT in [32, 64, 128, 256] + for num_warps in [2, 4, 8] + for num_stages in [1, 2, 3, 4] + ], + key=['B', 'H', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_global_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_nh = tl.program_id(0) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + T = eos - bos + + b_z = tl.zeros([], dtype=tl.float32) + NT = tl.cdiv(T, BT) + for i_c in range(NT): + i_t = NT - 1 - i_c if REVERSE else i_c + if HEAD_FIRST: + p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + else: + p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + b_ss = tl.sum(b_s, 0) + if REVERSE: + b_o = -b_o + b_ss + b_s + b_o += b_z + if i_c >= 0: + b_z += b_ss + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for BT in [16, 32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [1, 2, 3, 4] + ], + key=['B', 'H', 'S', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_global_cumsum_vector_kernel( + s, + o, + scale, + cu_seqlens, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + T = eos - bos + + b_z = tl.zeros([BS], dtype=tl.float32) + NT = tl.cdiv(T, BT) + for i_c in range(NT): + i_t = NT - 1 - i_c if REVERSE else i_c + if HEAD_FIRST: + p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + else: + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if REVERSE: + b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) + else: + b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) + if HAS_SCALE: + b_c *= scale + tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + b_z += tl.sum(b_s, 0) + + +def chunk_local_cumsum_scalar( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + if head_first: + B, H, T = g.shape + else: + B, T, H = g.shape + assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = (NT, B * H) + chunk_local_cumsum_scalar_kernel[grid]( + s=g_org, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g + + +def chunk_local_cumsum_vector( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + if head_first: + B, H, T, S = g.shape + else: + B, T, H, S = g.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" + + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) + # keep cummulative normalizer in fp32 + # this kernel is equivalent to + # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) + chunk_local_cumsum_vector_kernel[grid]( + s=g_org, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + S=S, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g + + +@input_guard +def chunk_global_cumsum_scalar( + s: torch.Tensor, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + scale: float = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T = s.shape + else: + B, T, H = s.shape + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + + z = torch.empty_like(s, dtype=output_dtype or s.dtype) + grid = (N * H,) + chunk_global_cumsum_scalar_kernel[grid]( + s=s, + o=z, + scale=scale, + cu_seqlens=cu_seqlens, + T=T, + B=B, + H=H, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return z + + +@input_guard +def chunk_global_cumsum_vector( + s: torch.Tensor, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + scale: float = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T, S = s.shape + else: + B, T, H, S = s.shape + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BS = min(32, triton.next_power_of_2(S)) + + z = torch.empty_like(s, dtype=output_dtype or s.dtype) + grid = (triton.cdiv(S, BS), N * H) + chunk_global_cumsum_vector_kernel[grid]( + s=s, + o=z, + scale=scale, + cu_seqlens=cu_seqlens, + T=T, + B=B, + H=H, + S=S, + BS=BS, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return z + + +@input_guard +def chunk_global_cumsum( + s: torch.Tensor, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + scale: float = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + if cu_seqlens is not None: + assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" + if len(s.shape) == 3: + return chunk_global_cumsum_scalar( + s=s, + reverse=reverse, + cu_seqlens=cu_seqlens, + scale=scale, + head_first=head_first, + output_dtype=output_dtype, + ) + elif len(s.shape) == 4: + return chunk_global_cumsum_vector( + s=s, + reverse=reverse, + cu_seqlens=cu_seqlens, + scale=scale, + head_first=head_first, + output_dtype=output_dtype, + ) + else: + raise ValueError( + f"Unsupported input shape {s.shape}, " + f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " + f"or [B, H, T]/[B, H, T, D] otherwise", + ) + + +@input_guard +def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, + chunk_indices: torch.LongTensor | None = None, + **kwargs, +) -> torch.Tensor: + if cu_seqlens is not None: + assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" + if len(g.shape) == 3: + return chunk_local_cumsum_scalar( + g=g, + chunk_size=chunk_size, + reverse=reverse, + scale=scale, + cu_seqlens=cu_seqlens, + head_first=head_first, + output_dtype=output_dtype, + chunk_indices=chunk_indices, + ) + elif len(g.shape) == 4: + return chunk_local_cumsum_vector( + g=g, + chunk_size=chunk_size, + reverse=reverse, + scale=scale, + cu_seqlens=cu_seqlens, + head_first=head_first, + output_dtype=output_dtype, + chunk_indices=chunk_indices, + ) + else: + raise ValueError( + f"Unsupported input shape {g.shape}, " + f"which should be (B, T, H, D) if `head_first=False` " + f"or (B, H, T, D) otherwise", + ) diff --git a/fla/ops/utils/index.py b/fla/ops/utils/index.py new file mode 100644 index 0000000000000000000000000000000000000000..e7bf3ed5fe4a309bc0a2c1c7c8f0a22abee36d7c --- /dev/null +++ b/fla/ops/utils/index.py @@ -0,0 +1,141 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.utils import autotune_cache_kwargs, tensor_cache + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [4, 8, 16, 32] + ], + key=['B'], + **autotune_cache_kwargs, +) +@triton.jit +def prepare_position_ids_kernel( + y, + cu_seqlens, + B: tl.constexpr, +): + i_n = tl.program_id(0) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + + o = tl.arange(0, B) + for i in range(0, tl.cdiv(T, B) * B, B): + o_i = o + i + tl.store(y + bos + o_i, o_i, o_i < T) + + +@tensor_cache +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return torch.diff(cu_seqlens) + + +@tensor_cache +def prepare_lens_from_mask(mask: torch.BoolTensor) -> torch.LongTensor: + return mask.sum(dim=-1, dtype=torch.int32) + + +@tensor_cache +def prepare_cu_seqlens_from_lens( + lens: torch.LongTensor, + dtype: torch.dtype | None = torch.int32, +) -> torch.LongTensor: + return F.pad(lens.cumsum(dim=0, dtype=dtype), (1, 0)) + + +@tensor_cache +def prepare_cu_seqlens_from_mask( + mask: torch.BoolTensor, + dtype: torch.dtype | None = torch.int32, +) -> torch.LongTensor: + return prepare_cu_seqlens_from_lens(prepare_lens_from_mask(mask), dtype) + + +@tensor_cache +def prepare_split_cu_seqlens( + batch_size: int, + seq_len: int, + split_size: int, + cu_seqlens: torch.LongTensor | None = None, + dtype: torch.dtype | None = torch.int32, + device: torch.device | None = torch.device('cpu'), +) -> torch.LongTensor: + if cu_seqlens is None: + total_tokens = batch_size * seq_len + cu_seqlens = list(range(0, total_tokens, seq_len)) + [total_tokens] + else: + cu_seqlens = cu_seqlens.tolist() + return torch.tensor( + [ + i + for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) + for i in range(bos, eos, split_size) + ] + [cu_seqlens[-1]], + dtype=dtype, + device=device, + ) + + +@tensor_cache +def prepare_position_ids(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: + if cu_seqlens_cpu is not None: + return torch.cat([ + torch.arange(n, dtype=cu_seqlens.dtype, device=cu_seqlens.device) + for n in prepare_lens(cu_seqlens_cpu).unbind() + ]) + return torch.cat([ + torch.arange(n, dtype=cu_seqlens.dtype, device=cu_seqlens.device) + for n in prepare_lens(cu_seqlens).unbind() + ]) + + +@tensor_cache +def prepare_sequence_ids(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: + return prepare_position_ids(cu_seqlens, cu_seqlens_cpu).eq(0).cumsum(0) - 1 + + +@tensor_cache +def prepare_token_indices(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: + position_ids = prepare_position_ids(cu_seqlens, cu_seqlens_cpu) + return torch.stack([prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu), position_ids], 1).to(cu_seqlens) + + +@tensor_cache +def prepare_chunk_indices( + cu_seqlens: torch.LongTensor, + chunk_size: int, + cu_seqlens_cpu: torch.LongTensor | None = None, +) -> torch.LongTensor: + if cu_seqlens_cpu is not None: + indices = torch.cat([torch.arange(n, device=cu_seqlens.device) + for n in triton.cdiv(prepare_lens(cu_seqlens_cpu), chunk_size).tolist()]) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) + indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) + + +@tensor_cache +def prepare_chunk_offsets( + cu_seqlens: torch.LongTensor, + chunk_size: int, +) -> torch.LongTensor: + return F.pad(triton.cdiv(prepare_lens(cu_seqlens), chunk_size), (1, 0), value=0).cumsum(-1) + + +@tensor_cache +def get_max_num_splits( + cu_seqlens: torch.LongTensor, + chunk_size: int, + cu_seqlens_cpu: torch.LongTensor | None = None +) -> int: + if cu_seqlens_cpu is not None: + return triton.cdiv(int(max(prepare_lens(cu_seqlens_cpu))), chunk_size) + return triton.cdiv(int(max(prepare_lens(cu_seqlens))), chunk_size) diff --git a/fla/ops/utils/logcumsumexp.py b/fla/ops/utils/logcumsumexp.py new file mode 100644 index 0000000000000000000000000000000000000000..b7c19de6748164e5342043ee4010444a206e88d6 --- /dev/null +++ b/fla/ops/utils/logcumsumexp.py @@ -0,0 +1,53 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import autotune_cache_kwargs + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for BT in [16, 32, 64] + for num_warps in [2, 4, 8] + ], + key=['S'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def logcumsumexp_fwd_kernel( + s, + z, + T, + S: tl.constexpr, + BT: tl.constexpr, +): + i_bh = tl.program_id(0) + o_i = tl.arange(0, BT) + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) + + b_mp = tl.full([S], float('-inf'), dtype=tl.float32) + b_zp = tl.zeros([S], dtype=tl.float32) + for i_t in range(tl.cdiv(T, BT)): + p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) + + # [BT, S] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + # [S,] + b_mc = tl.max(b_s, 0) + b_mc = tl.maximum(b_mp, b_mc) + b_zp = b_zp * exp(b_mp - b_mc) + # [BT, S] + b_s = exp(b_s - b_mc) + b_z = tl.dot(m_s, b_s, allow_tf32=False) + b_zp + # [S,] + b_zc = tl.max(b_z, 0) + b_mp = b_mc + b_zp = b_zc + # [BT, BS] + # small eps to prevent underflows + b_z = log(tl.where(b_z != 0, b_z, 1e-20)) + b_mc + tl.store(p_z, b_z.to(p_z.dtype.element_ty), boundary_check=(0, 1)) diff --git a/fla/ops/utils/logsumexp.py b/fla/ops/utils/logsumexp.py new file mode 100644 index 0000000000000000000000000000000000000000..f8f0e664a06f6608a4f0b39ad9cb63ffae1801b2 --- /dev/null +++ b/fla/ops/utils/logsumexp.py @@ -0,0 +1,80 @@ +# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit +def logsumexp_fwd_kernel( + x, + z, + scale, + D: tl.constexpr, + B: tl.constexpr, + HAS_SCALE: tl.constexpr, +): + i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + o_d = i_d * B + tl.arange(0, B) + m_d = o_d < D + + b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf')) + if HAS_SCALE: + b_x = b_x * scale + b_m = tl.max(b_x, 0) + b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m + tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z) + + +def logsumexp_fwd( + x, + scale: float | None = None, + dtype: torch.dtype | None = None, +): + r""" + Compute the logsumexp of the input tensor over the last dimension. + + Args: + x (Tensor): + The input tensor of any shape. + scale (Optional[float]): + The scale applied to the input tensor. Default: `None`. + dtype (Optional[torch.dtype]): + The data type of the output tensor. Default: `None`. + Returns: + Tensor: The logsumexp of the input tensor. + """ + + shape = x.shape + x = x.view(-1, shape[-1]) + N, D = x.shape + B = min(triton.next_power_of_2(D), 64 * 1024) + ND = triton.cdiv(D, B) + + z = x.new_empty(N, ND, dtype=torch.float) + logsumexp_fwd_kernel[(N, ND)]( + x=x, + z=z, + scale=scale, + D=D, + B=B, + ) + z = z.logsumexp(-1).view(*shape[:-1]) + if dtype is not None and dtype != torch.float: + z = z.to(dtype) + return z diff --git a/fla/ops/utils/matmul.py b/fla/ops/utils/matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..b565f8c0efb5f32c4978fe14ec548c4bed62572a --- /dev/null +++ b/fla/ops/utils/matmul.py @@ -0,0 +1,244 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# code adapted from +# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, input_guard + + +# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes: +# - A list of `triton.Config` objects that define different configurations of +# meta-parameters (e.g., `BM`) and compilation options (e.g., `num_warps`) to try +# - An auto-tuning *key* whose change in values will trigger evaluation of all the +# provided configs +@triton.heuristics({ + 'HAS_ALPHA': lambda args: args['alpha'] is not None, + 'HAS_BETA': lambda args: args['beta'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BM': 128, 'BK': 64, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8), + triton.Config({'BM': 64, 'BK': 32, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4), + triton.Config({'BM': 128, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), + triton.Config({'BM': 128, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), + triton.Config({'BM': 64, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), + triton.Config({'BM': 128, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4), + triton.Config({'BM': 64, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=5, num_warps=2), + triton.Config({'BM': 32, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=5, num_warps=2), + # Good config for fp8 inputs. + # triton.Config({'BM': 128, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8), + # triton.Config({'BM': 256, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=3, num_warps=8), + # triton.Config({'BM': 256, 'BK': 128, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), + # triton.Config({'BM': 64, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4), + # triton.Config({'BM': 128, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), + # triton.Config({'BM': 128, 'BK': 64, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), + # triton.Config({'BM': 64, 'BK': 64, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), + # triton.Config({'BM': 128, 'BK': 64, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4) + ], + key=['M', 'N', 'K'], + **autotune_cache_kwargs, +) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a, + b, + c, + input, + alpha, + beta, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `s_am` is how much to increase `a` + # by to get the element one row down (A has M rows). + stride_ab, stride_am, stride_ak, # a: batch, M, K + stride_bk, stride_bn, # b: K, N + stride_cb, stride_cm, stride_cn, # c: batch, M, N + # Meta-parameters + BM: tl.constexpr, + BK: tl.constexpr, + BN: tl.constexpr, + G: tl.constexpr, + ACTIVATION: tl.constexpr, + HAS_INPUT: tl.constexpr, + HAS_ALPHA: tl.constexpr, + HAS_BETA: tl.constexpr, + ALLOW_TF32: tl.constexpr, + X_DIM: tl.constexpr = 1, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + i_b, i_m, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + NM, NN = tl.num_programs(1), tl.num_programs(2) + i_m, i_n = tl.swizzle2d(i_m, i_n, NM, NN, G) + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `p_a` is a block of [BM, BK] pointers + # `p_b` is a block of [BK, BN] pointers + # See above `Pointer Arithmetic` section for details + a_batch_ptr = a + i_b * stride_ab + o_am = (i_m * BM + tl.arange(0, BM)) % M + o_bn = (i_n * BN + tl.arange(0, BN)) % N + o_k = tl.arange(0, BK) + + p_a = a_batch_ptr + (o_am[:, None] * stride_am + o_k[None, :] * stride_ak) + p_b = b + (o_k[:, None] * stride_bk + o_bn[None, :] * stride_bn) + + b_acc = tl.zeros((BM, BN), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BK)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + b_a = tl.load(p_a, mask=o_k[None, :] < K - k * BK, other=0.0) + b_b = tl.load(p_b, mask=o_k[:, None] < K - k * BK, other=0.0) + # We accumulate along the K dimension. + b_acc = tl.dot(b_a, b_b, acc=b_acc, allow_tf32=ALLOW_TF32) + # Advance the ptrs to the next K block. + p_a += BK * stride_ak + p_b += BK * stride_bk + + o_cm = i_m * BM + tl.arange(0, BM) + o_cn = i_n * BN + tl.arange(0, BN) + mask = (o_cm[:, None] < M) & (o_cn[None, :] < N) + + b_c = b_acc + # You can fuse arbitrary activation functions here + # while the b_acc is still in FP32! + if ACTIVATION == "leaky_relu": + b_c = leaky_relu(b_c) + elif ACTIVATION == "relu": + b_c = relu(b_c) + elif ACTIVATION == "sigmoid": + b_c = sigmoid(b_c) + elif ACTIVATION == "tanh": + b_c = tanh(b_c) + + if HAS_ALPHA: + b_c *= tl.load(alpha) + + if HAS_INPUT: + p_i = input + (stride_cm * o_cm[:, None] if X_DIM == 2 else 0) + stride_cn * o_cn[None, :] + mask_p = (o_cn[None, :] < N) if X_DIM == 1 else mask + b_i = tl.load(p_i, mask=mask_p, other=0.0).to(tl.float32) + if HAS_BETA: + b_i *= tl.load(beta) + b_c += b_i + + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + c_batch_ptr = c + i_b * stride_cb + p_c = c_batch_ptr + stride_cm * o_cm[:, None] + stride_cn * o_cn[None, :] + tl.store(p_c, b_c.to(c.dtype.element_ty), mask=mask) + + +# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`. +@triton.jit +def leaky_relu(x): + return tl.where(x >= 0, x, 0.01 * x) + + +@triton.jit +def sigmoid(x): + # σ(x) = 1 / (1 + exp(-x)) + return 1.0 / (1.0 + exp(-x)) + + +@triton.jit +def tanh(x): + # tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)) + # 2 * sigmoid(2x) - 1 + return (exp(x) - exp(-x)) / (exp(x) + exp(-x)) + + +@triton.jit +def relu(x): + # ReLU(x) = max(0, x) + return tl.maximum(x, 0.0) + + +@input_guard +def matmul(a, b, activation=''): + assert a.dim() in [2, 3], "a must be 2D or 3D" + assert b.dim() == 2, "b must be 2D" + assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}" + + if a.dim() == 2: + a_dim = 2 + a = a.unsqueeze(0).contiguous() # (1, M, K) + else: + a_dim = 3 + allow_tf32 = False if a.dtype == torch.float32 else True + + B, M, K = a.shape[0], a.shape[1], a.shape[2] + K_b, N = b.shape + assert K_b == K, f"Incompatible K dimension: A {K} vs B {K_b}" + c = a.new_empty(B, M, N) + + def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN'])) + matmul_kernel[grid]( + a, b, c, None, None, None, + M, N, K, + a.stride(0), a.stride(1), a.stride(2), # stride_ab, stride_am, stride_ak + b.stride(0), b.stride(1), # stride_bk, stride_bn (b.dim() == 2) + c.stride(0), c.stride(1), c.stride(2), # stride_cb, stride_cm, stride_cn + ACTIVATION=activation, + ALLOW_TF32=allow_tf32, + HAS_INPUT=False, + ) + return c.squeeze(0) if a_dim == 2 else c + + +@input_guard +def addmm( + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + alpha: float | None = None, + beta: float | None = None, +) -> torch.Tensor: + assert a.dim() in [2, 3], "a must be 2D or 3D" + assert b.dim() == 2, "b must be 2D" + assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}" + + if a.dim() == 2: + a_dim = 2 + a = a.unsqueeze(0).contiguous() # (1, M, K) + else: + a_dim = 3 + allow_tf32 = False if a.dtype == torch.float32 else True + + B, M, K = a.shape[0], a.shape[1], a.shape[2] + K_b, N = b.shape + assert K_b == K, f"Incompatible K dimension: A {K} vs B {K_b}" + c = a.new_empty(B, M, N) + + def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN'])) + matmul_kernel[grid]( + a, b, c, x, alpha, beta, + M, N, K, + a.stride(0), a.stride(1), a.stride(2), # stride_ab, stride_am, stride_ak + b.stride(0), b.stride(1), # stride_bk, stride_bn (b.dim() == 2) + c.stride(0), c.stride(1), c.stride(2), # stride_cb, stride_cm, stride_cn + ACTIVATION=None, + ALLOW_TF32=allow_tf32, + HAS_INPUT=True, + X_DIM=x.dim(), + ) + return c.squeeze(0) if a_dim == 2 else c diff --git a/fla/ops/utils/op.py b/fla/ops/utils/op.py new file mode 100644 index 0000000000000000000000000000000000000000..dd64d798e6aa384e8c6e576e7f6d2fb39adfcf55 --- /dev/null +++ b/fla/ops/utils/op.py @@ -0,0 +1,64 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import os + +import triton +import triton.language as tl +import triton.language.extra.libdevice as tldevice + +from fla.utils import IS_GATHER_SUPPORTED + +if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': + @triton.jit + def exp(x): return tldevice.fast_expf(x.to(tl.float32)) + @triton.jit + def exp2(x): return tldevice.exp2(x.to(tl.float32)) + @triton.jit + def log(x): return tldevice.fast_logf(x.to(tl.float32)) + @triton.jit + def log2(x): return tldevice.fast_log2f(x.to(tl.float32)) +else: + @triton.jit + def exp(x): return tl.exp(x.to(tl.float32)) + @triton.jit + def exp2(x): return tl.math.exp2(x.to(tl.float32)) + @triton.jit + def log(x): return tl.log(x.to(tl.float32)) + @triton.jit + def log2(x): return tl.log2(x.to(tl.float32)) + + +if not IS_GATHER_SUPPORTED: + @triton.jit + def gather(src, index, axis, _builder=None): + """ + Gather operation that works when tl.gather is not supported. + This is a fallback implementation that returns None. + Just to make triton compiler happy. + """ + return None +else: + gather = tl.gather + + +if hasattr(triton.language, '_experimental_make_tensor_descriptor'): + # For Triton 3.3.x + make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor +elif hasattr(triton.language, 'make_tensor_descriptor'): + # For Triton 3.4.x and later + make_tensor_descriptor = triton.language.make_tensor_descriptor +else: + """ + Fallback implementation when TMA is not supported. + Returns None to indicate TMA descriptors are unavailable. + Just make triton compiler happy. + """ + @triton.jit + def make_tensor_descriptor( + base, + shape, + strides, + block_shape, + _builder=None, + ): + return None diff --git a/fla/ops/utils/pack.py b/fla/ops/utils/pack.py new file mode 100644 index 0000000000000000000000000000000000000000..6bb14607a39a3e124d75ecec985a057cc189cc92 --- /dev/null +++ b/fla/ops/utils/pack.py @@ -0,0 +1,207 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Code adapted from https://github.com/mayank31398/cute-kernels + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_lens +from fla.utils import autotune_cache_kwargs, input_guard + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [4, 8, 16, 32] + ], + key=['D', 'PADDING_SIDE', 'PACK'], + **autotune_cache_kwargs, +) +@triton.jit +def packunpack_sequence_kernel( + x, + y, + cu_seqlens, + S, + D, + BD: tl.constexpr, + PADDING_SIDE: tl.constexpr, + PACK: tl.constexpr, +): + i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1) + + T = eos - bos + if PADDING_SIDE == 'left': + NP = S - T + if i_s < NP: + return + i_t = bos + (i_s - NP) + else: + if i_s >= T: + return + i_t = bos + i_s + + o_d = i_d * BD + tl.arange(0, BD) + mask = o_d < D + + if PACK: + b_x = tl.load(x + (i_b * S + i_s) * D + o_d, mask=mask) + tl.store(y + i_t * D + o_d, b_x, mask=mask) + else: + b_x = tl.load(x + i_t * D + o_d, mask=mask) + tl.store(y + (i_b * S + i_s) * D + o_d, b_x, mask=mask) + + +def pack_sequence_fwdbwd( + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str, +) -> torch.Tensor: + B, S = x.shape[:2] + D = x.numel() // (B * S) + BD = min(triton.next_power_of_2(D), 4096) + ND = triton.cdiv(D, BD) + + y = torch.empty(cu_seqlens[-1].item(), *x.shape[2:], device=x.device, dtype=x.dtype) + packunpack_sequence_kernel[ND, S, B]( + x=x, + y=y, + cu_seqlens=cu_seqlens, + S=S, + D=D, + BD=BD, + PADDING_SIDE=padding_side, + PACK=True, + ) + return y + + +def unpack_sequence_fwdbwd( + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str, + desired_shape: torch.Size, +) -> torch.Tensor: + if desired_shape is None: + desired_shape = (len(cu_seqlens) - 1, prepare_lens(cu_seqlens).max().item(), *x.shape[1:]) + y = torch.zeros(desired_shape, device=x.device, dtype=x.dtype) + B, S = y.shape[:2] + D = y.numel() // (B * S) + BD = min(triton.next_power_of_2(D), 4096) + ND = triton.cdiv(D, BD) + + packunpack_sequence_kernel[ND, S, B]( + x=x, + y=y, + cu_seqlens=cu_seqlens, + S=S, + D=D, + BD=BD, + PADDING_SIDE=padding_side, + PACK=False, + ) + return y + + +class PackSequenceFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str, + ) -> torch.Tensor: + assert padding_side in ['left', 'right'] + assert x.ndim >= 2 + + ctx.cu_seqlens = cu_seqlens + ctx.padding_side = padding_side + ctx.desired_shape = x.shape + + y = pack_sequence_fwdbwd( + x=x, + cu_seqlens=cu_seqlens, + padding_side=padding_side, + ) + return y + + @staticmethod + @input_guard + def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]: + dx = unpack_sequence_fwdbwd( + x=dy, + cu_seqlens=ctx.cu_seqlens, + padding_side=ctx.padding_side, + desired_shape=ctx.desired_shape, + ) + return dx, *[None] * 10 + + +class UnpackSequenceFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str, + desired_shape: torch.Size | None = None, + ) -> torch.Tensor: + assert padding_side in ['left', 'right'] + assert x.ndim >= 2 + if desired_shape is not None: + assert desired_shape[0] == cu_seqlens.shape[0] - 1 + assert desired_shape[2:] == x.shape[1:] + + ctx.cu_seqlens = cu_seqlens + ctx.padding_side = padding_side + + y = unpack_sequence_fwdbwd( + x=x, + cu_seqlens=cu_seqlens, + padding_side=padding_side, + desired_shape=desired_shape, + ) + return y + + @staticmethod + @input_guard + def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]: + dx = pack_sequence_fwdbwd( + x=dy, + cu_seqlens=ctx.cu_seqlens, + padding_side=ctx.padding_side, + ) + return dx, None, None, None + + +def pack_sequence( + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str = 'left', +) -> torch.Tensor: + return PackSequenceFunction.apply( + x, + cu_seqlens, + padding_side, + ) + + +def unpack_sequence( + x: torch.Tensor, + cu_seqlens: torch.Tensor, + padding_side: str = 'left', + desired_shape: torch.Size | None = None, +) -> torch.Tensor: + return UnpackSequenceFunction.apply( + x, + cu_seqlens, + padding_side, + desired_shape, + ) diff --git a/fla/ops/utils/pooling.py b/fla/ops/utils/pooling.py new file mode 100644 index 0000000000000000000000000000000000000000..76773c8a92328eabc5121ea8950de1111f046563 --- /dev/null +++ b/fla/ops/utils/pooling.py @@ -0,0 +1,211 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [16, 32, 64, 128] + for num_warps in [1, 2, 4, 8] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def mean_pooling_fwd_kernel( + x, + o, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) + # [BT, BD] + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + # [BD] + b_o = tl.sum(b_x, axis=0) / min(BT, T - i_t * BT) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [16, 32, 64, 128] + for num_warps in [1, 2, 4, 8] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def mean_pooling_bwd_kernel( + do, + dx, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) + # [BD] + b_do = tl.load(p_do, boundary_check=(0,)).to(tl.float32) + # [BT, BD] + b_dx = b_do / tl.full((BT,), min(BT, T - i_t * BT), dtype=tl.float32)[:, None] + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + + +def mean_pooling_fwd( + x: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, D = x.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = x.new_empty(B, NT, H, D) + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) + mean_pooling_fwd_kernel[grid]( + x, + o, + cu_seqlens, + chunk_indices, + T=T, + H=H, + D=D, + BT=BT, + ) + return o + + +def mean_pooling_bwd( + do: torch.Tensor, + batch_size: int, + seq_len: int, + chunk_size: int, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> torch.Tensor: + B, T, H, D = batch_size, seq_len, *do.shape[-2:] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dx = do.new_empty(B, T, H, D) + def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) + mean_pooling_bwd_kernel[grid]( + do, + dx, + cu_seqlens, + chunk_indices, + T=T, + H=H, + D=D, + BT=BT, + ) + return dx + + +class MeanPoolingFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + x: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.LongTensor | None = None, + ) -> torch.Tensor: + o = mean_pooling_fwd(x, chunk_size, cu_seqlens) + ctx.batch_size = x.shape[0] + ctx.seq_len = x.shape[1] + ctx.chunk_size = chunk_size + ctx.cu_seqlens = cu_seqlens + return o + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, do, + ) -> tuple[torch.Tensor, None, None]: + batch_size = ctx.batch_size + seq_len = ctx.seq_len + chunk_size = ctx.chunk_size + cu_seqlens = ctx.cu_seqlens + dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, cu_seqlens) + return dx, None, None + + +def mean_pooling( + x: torch.Tensor, + chunk_size: int, + cu_seqlens: torch.LongTensor | None = None, + head_first: bool = False, +) -> torch.Tensor: + if head_first: + x = x.transpose(1, 2) + if cu_seqlens is not None: + if x.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {x.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + o = MeanPoolingFunction.apply(x, chunk_size, cu_seqlens) + if head_first: + o = o.transpose(1, 2) + return o diff --git a/fla/ops/utils/softmax.py b/fla/ops/utils/softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..b0988f1b9291016259bc6636d755ad8adfe2c8d3 --- /dev/null +++ b/fla/ops/utils/softmax.py @@ -0,0 +1,105 @@ +# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.utils import IS_AMD, autotune_cache_kwargs + +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit +def softmax_fwd_kernel( + x, + p, + D: tl.constexpr, + B: tl.constexpr, +): + i_n = tl.program_id(0) + o_d = tl.arange(0, B) + m_d = o_d < D + + b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf')) + b_m = tl.max(b_x, 0) + b_x = exp(b_x - b_m) + b_p = b_x / tl.sum(b_x, 0) + + tl.store(p + i_n * D + o_d, b_p.to(p.dtype.element_ty), mask=m_d) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D'], + **autotune_cache_kwargs, +) +@triton.jit +def softmax_bwd_kernel( + p, + dp, + ds, + D: tl.constexpr, + B: tl.constexpr, +): + i_n = tl.program_id(0) + o_d = tl.arange(0, B) + m_d = o_d < D + + b_p = tl.load(p + i_n * D + o_d, mask=m_d, other=0.) + b_dp = tl.load(dp + i_n * D + o_d, mask=m_d, other=0.) + b_pp = tl.sum(b_p * b_dp, 0) + b_ds = b_p * b_dp - b_p * b_pp + tl.store(ds + i_n * D + o_d, b_ds.to(ds.dtype.element_ty), mask=m_d) + + +def softmax_fwd( + x: torch.Tensor, + dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + shape = x.shape + x = x.view(-1, x.shape[-1]) + + N, D = x.shape + B = triton.next_power_of_2(D) + + p = torch.empty_like(x, dtype=dtype) + softmax_fwd_kernel[(N,)]( + x=x, + p=p, + D=D, + B=B, + ) + return p.view(*shape) + + +def softmax_bwd( + p: torch.Tensor, + dp: torch.Tensor, + dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + shape = p.shape + p = p.view(-1, p.shape[-1]) + ds = torch.empty_like(p, dtype=dtype) + + N, D = p.shape + B = triton.next_power_of_2(D) + softmax_bwd_kernel[(N,)]( + p=p, + dp=dp, + ds=ds, + D=D, + B=B, + ) + return ds.view(*shape) diff --git a/fla/ops/utils/softplus.py b/fla/ops/utils/softplus.py new file mode 100644 index 0000000000000000000000000000000000000000..225f865a7a0ccf1012057e335ef1f21edd0ada35 --- /dev/null +++ b/fla/ops/utils/softplus.py @@ -0,0 +1,104 @@ +# REVISED FROM +# https://github.com/shawntan/stickbreaking-attention/blob/main/stickbreaking_attention/sb_varlen/softplus.py + +import triton +from triton import language as tl + +from fla.utils import IS_NVIDIA + + +def _generate_softplus(num_pack): + template = """ + .reg .pred p; + setp.gt.f32 p, ${in_reg}, 20.; + @p mov.f32 ${out_reg}, ${in_reg}; + @!p mul.f32 ${out_reg}, ${in_reg}, 1.4426950408889634; + @!p ex2.approx.ftz.f32 ${out_reg}, ${out_reg}; + @!p add.f32 ${out_reg}, ${out_reg}, 1.0; + @!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg}; + @!p mul.f32 ${out_reg}, ${out_reg}, 0.6931471805599453; + """ + out_str = "" + + for i in range(num_pack): + inner_str = template.format(out_reg=i, in_reg=i + num_pack) + out_str += "{" + inner_str + "}\n" + # flatten out because torch.compile doesn't like newlines + out_str = " ".join(out_str.split("\n")) + return out_str + + +def _generate_softplus2(num_pack): + template = """ + .reg .pred p; + setp.gt.f32 p, ${in_reg}, 15.; + @p mov.f32 ${out_reg}, ${in_reg}; + @!p ex2.approx.ftz.f32 ${out_reg}, ${in_reg}; + @!p add.f32 ${out_reg}, ${out_reg}, 1.0; + @!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg}; + """ + out_str = "" + + for i in range(num_pack): + inner_str = template.format(out_reg=i, in_reg=i + num_pack) + out_str += "{" + inner_str + "}\n" + # flatten out because torch.compile doesn't like newlines + out_str = " ".join(out_str.split("\n")) + return out_str + + +def _generate_constraints(num_pack): + return ",".join("=r" for i in range(num_pack)) + "," + ",".join("r" for i in range(num_pack)) + + +_NUM_REG = 1 +s_softplus: tl.constexpr = tl.constexpr(_generate_softplus(_NUM_REG)) +s_softplus2: tl.constexpr = tl.constexpr(_generate_softplus2(_NUM_REG)) +s_constraints: tl.constexpr = tl.constexpr(_generate_constraints(_NUM_REG)) +NUM_REG: tl.constexpr = tl.constexpr(_NUM_REG) + + +@triton.jit +def softplus_nv(x): + # equivalent to: + # return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x) + return tl.inline_asm_elementwise( + asm=s_softplus, + constraints=s_constraints, + pack=NUM_REG, + args=[ + x, + ], + dtype=tl.float32, + is_pure=True, + ) + +@triton.jit +def softplus_triton(x): + return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x) + +@triton.jit +def softplus2_nv(x): + # equivalent to: + # return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x) + return tl.inline_asm_elementwise( + asm=s_softplus2, + constraints=s_constraints, + pack=NUM_REG, + args=[ + x, + ], + dtype=tl.float32, + is_pure=True, + ) + +@triton.jit +def softplus2_triton(x): + return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x) + +if IS_NVIDIA: + softplus = softplus_nv + softplus2 = softplus2_nv +else: + softplus = softplus_triton + softplus2 = softplus2_triton diff --git a/fla/ops/utils/solve_tril.py b/fla/ops/utils/solve_tril.py new file mode 100644 index 0000000000000000000000000000000000000000..f13dfb8d53d64619ced75baf0f7f5efc2552a610 --- /dev/null +++ b/fla/ops/utils/solve_tril.py @@ -0,0 +1,392 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import os + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.op import make_tensor_descriptor +from fla.utils import IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard + +FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') +assert FLA_TRIL_PRECISION in ['ieee', 'tf32', 'tf32x3'], \ + f"FLA_TRIL_PRECISION must be one of 'ieee', 'tf32', or 'tf32x3', but got {FLA_TRIL_PRECISION}" +DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION}) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'DOT_PRECISION': 'ieee'}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4, 5] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def solve_tril_16x16_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + A = A + (bos*H + i_h) * BT + Ai = Ai + (bos*H + i_h) * 16 + + offset = (i_t * 16) % BT + if not USE_TMA: + p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * 16, offset), (16, 16), (1, 0)) + # [16, 16] + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = tl.where(m_A, b_A, 0) + else: + desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, 16], [H*16, 1], [16, 16]) + b_A = desc.load([i_t * 16, offset]).to(tl.float32) + b_A = tl.where(m_A, b_A, 0) + b_A = -b_A + + for i in range(2, min(16, T - i_t * 16)): + # [16] + b_a = -tl.load(A + (i_t * 16 + i) * H*BT + o_i + offset) + b_a = tl.where(o_i < i, b_a, 0.) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) + b_A = tl.where((o_i == i)[:, None], b_a, b_A) + b_A += m_I + if not USE_TMA: + p_Ai = tl.make_block_ptr(Ai, (T, 16), (H*16, 1), (i_t * 16, 0), (16, 16), (1, 0)) + tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + else: + desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4, 5] + for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST + ], + key=['H', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + + b_Ai_11 += m_I + b_Ai_22 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + + b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + else: + desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4, 5] + for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST + ], + key=['H', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_A_33 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_A_44 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) + b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) + b_a_11 = tl.where(o_i < i, b_a_11, 0.) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) + b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(32 + 2, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 32) + b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(48 + 2, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 48) + b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_A_31 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_A_32 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_A_41 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_A_42 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_A_43 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) + b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) + b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) + b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) + b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) + + b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) + b_Ai_32 = -tl.dot(tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), b_Ai_22, input_precision=DOT_PRECISION) + b_Ai_43 = -tl.dot(tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), b_Ai_33, input_precision=DOT_PRECISION) + + b_Ai_31 = -tl.dot( + b_Ai_33, + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_42 = -tl.dot( + b_Ai_44, + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_41 = -tl.dot( + b_Ai_44, + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) + p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) + p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) + p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) + p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) + p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) + p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) + p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) + tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + else: + desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")) + desc_o.store([i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")) + + +@input_guard +def solve_tril( + A: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.LongTensor | None = None, + output_dtype: torch.dtype = torch.float, +) -> torch.Tensor: + """ + Compute the inverse of the matrix I + A + A should be strictly lower triangular, i.e., A.triu() == 0. + + Args: + A (torch.Tensor): + [B, T, H, BT], where BT should only be 16, 32, or 64. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. Default: `None`. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float`. + If `None`, the output dtype will be the same as the input dtype. + + Returns: + (I + A)^-1 with the same shape as A + """ + assert A.shape[-1] in [16, 32, 64] + output_dtype = A.dtype if output_dtype is None else output_dtype + + B, T, H, BT = A.shape + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + + Ai = torch.zeros_like(A, dtype=output_dtype) + if BT == 16: + merge_fn = solve_tril_16x16_kernel + elif BT == 32: + merge_fn = merge_16x16_to_32x32_inverse_kernel + elif BT == 64: + merge_fn = merge_16x16_to_64x64_inverse_kernel + + merge_fn[NT, B * H]( + A=A, + Ai=Ai, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + USE_TMA=IS_TMA_SUPPORTED, + ) + return Ai diff --git a/fla/utils.py b/fla/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d956a0d7fad613a8695e510d012e1c43d5da3a44 --- /dev/null +++ b/fla/utils.py @@ -0,0 +1,553 @@ +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import contextlib +import functools +import inspect +import logging +import os +import sys +import warnings +from collections.abc import Callable +from enum import Enum +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +import torch +import triton +from packaging import version + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from fla import __version__ + +FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" +FLA_CACHE_RESULTS = os.getenv('FLA_CACHE_RESULTS', '1') == '1' +FLA_DISABLE_TENSOR_CACHE = os.getenv('FLA_DISABLE_TENSOR_CACHE', '0') == '1' +TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0") +TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1") + + +SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters + +autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} + + +@lru_cache(maxsize=1) +def check_environments(): + """ + Checks the current operating system, Triton version, and Python version, + issuing warnings if they don't meet recommendations. + This function's body only runs once due to lru_cache. + """ + # Check Operating System + if sys.platform == 'win32': + # Check if triton-windows is installed + try: + from importlib.metadata import PackageNotFoundError, metadata + metadata('triton-windows') + # triton-windows is installed, no warning needed + except PackageNotFoundError: + logger.warning( + "Detected Windows operating system. Consider installing triton-windows " + "(https://github.com/triton-lang/triton-windows) for better compatibility. " + "Without it, some features may not work correctly.", + ) + + triton_version = version.parse(triton.__version__) + required_triton_version = version.parse("3.3.0") + + if triton_version < required_triton_version: + logger.warning( + f"Current Triton version {triton_version} is below the recommended 3.3.0 version. " + "Errors may occur and these issues will not be fixed. " + "Please consider upgrading Triton.", + ) + + # Check Python version + py_version = version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") + required_py_version = version.parse("3.11") + + if py_version < required_py_version: + logger.warning( + f"Current Python version {py_version} is below the recommended 3.11 version. " + "It is recommended to upgrade to Python 3.11 or higher for the best experience.", + ) + + return None + + +check_environments() + + +def get_abs_err(x, y): + return (x.detach()-y.detach()).flatten().abs().max().item() + + +def get_err_ratio(x, y): + err = (x.detach()-y.detach()).flatten().square().mean().sqrt().item() + base = (x.detach()).flatten().square().mean().sqrt().item() + return err / (base + 1e-8) + + +def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): + abs_atol = get_abs_err(ref, tri) + msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}" + logger.info(msg) + error_rate = get_err_ratio(ref, tri) + if abs_atol <= err_atol: + return + assert not torch.isnan(ref).any(), f"{prefix}: NaN detected in ref" + assert not torch.isnan(tri).any(), f"{prefix}: NaN detected in tri" + if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): + if error_rate > ratio: + warnings.warn(msg) + else: + assert error_rate < ratio, msg + + +def tensor_cache( + fn: Callable[..., torch.Tensor], +) -> Callable[..., torch.Tensor]: + """ + A decorator that caches the most recent result of a function with tensor inputs. + + This decorator will store the output of the decorated function for the most recent set of input tensors. + If the function is called again with the same input tensors, it will return the cached result. + + If FLA_DISABLE_TENSOR_CACHE environment variable is set to '1', caching is disabled. + + Args: + fn (Callable[..., torch.Tensor]): + The function to be decorated. It should take tensor inputs and return tensor outputs. + + Returns: + Callable[..., torch.Tensor]: + A wrapped version of the input function with single-entry caching. + """ + last_args: tuple | None = None + last_kwargs: dict | None = None + last_result: Any = None + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal last_args, last_kwargs, last_result + + # Skip cache if FLA_DISABLE_TENSOR_CACHE is set + if FLA_DISABLE_TENSOR_CACHE: + return fn(*args, **kwargs) + + if last_args is not None and last_kwargs is not None: + if len(args) == len(last_args) and len(kwargs) == len(last_kwargs): + if all(a is b for a, b in zip(args, last_args, strict=False)) and \ + all(k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items()): + return last_result + + result = fn(*args, **kwargs) + last_args, last_kwargs, last_result = args, kwargs, result + return result + + return wrapper + + +def input_guard( + fn: Callable[..., torch.Tensor] | None = None, + *, + no_guard_contiguous: bool | list[str] = False, +) -> Callable[[Callable[..., torch.Tensor]], Callable[..., torch.Tensor]] | Callable[..., torch.Tensor]: + """ + A decorator to make sure all input tensors are contiguous and set the device based on input tensors. + + Args: + no_guard_contiguous: If True, skip all contiguous checks. If a list of parameter names, skip contiguous check for those parameters. + """ + + def decorator(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + # Get function signature for parameter name mapping + sig = inspect.signature(fn) + param_names = list(sig.parameters.keys()) + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + # Convert no_guard_contiguous to list of parameter names if it's a list + skip_params = set() + if isinstance(no_guard_contiguous, list): + skip_params = set(no_guard_contiguous) + + # Process args with parameter name mapping + processed_args = [] + for i, arg in enumerate(args): + if i < len(param_names): + param_name = param_names[i] + else: + # For *args beyond signature, use position as name + param_name = f"__arg_{i}" + + if isinstance(arg, torch.Tensor): + if no_guard_contiguous is True or param_name in skip_params: + processed_args.append(arg) + else: + processed_args.append(arg.contiguous()) + else: + processed_args.append(arg) + + # Process kwargs + processed_kwargs = {} + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + if no_guard_contiguous is True or k in skip_params: + processed_kwargs[k] = v + else: + processed_kwargs[k] = v.contiguous() + else: + processed_kwargs[k] = v + + tensor = None + for arg in args: + if isinstance(arg, torch.Tensor): + tensor = arg + break + if tensor is None: + for value in kwargs.values(): + if isinstance(value, torch.Tensor): + tensor = value + break + + if tensor is not None: + ctx = custom_device_ctx(tensor.device.index) + else: + ctx = contextlib.nullcontext() + + with ctx: + return fn(*processed_args, **processed_kwargs) + + return wrapper + + # Handle direct usage without parentheses: @input_guard + if fn is not None: + return decorator(fn) + + return decorator + + +def contiguous(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """Alias for input_guard() without parameters.""" + return input_guard(fn) + + +def require_version(version, hint): + """ + Perform a runtime check of the dependency versions, using the exact same syntax used by pip. + """ + def decorator(fn): + @functools.wraps(fn) + def wrapper(ctx, *args, **kwargs): + from transformers.utils.versions import require_version + require_version(version, hint) + return fn(ctx, + *(i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args), + **{k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}) + return wrapper + return decorator + + +class Action(Enum): + NONE = "none" + NOTIFY = "notify" + NOTIFY_ALWAYS = "notify_always" + RAISE = "raise" + + +def deprecate_kwarg( + old_name: str, + version: str, + new_name: str | None = None, + warn_if_greater_or_equal_version: bool = False, + raise_if_greater_or_equal_version: bool = False, + raise_if_both_names: bool = False, + additional_message: str | None = None, +): + """ + Decorator to notify users about deprecated keyword arguments, replacing them with a new name if specified. + + This decorator allows you to: + - Notify users when a keyword argument is deprecated. + - Automatically replace deprecated keyword arguments with new ones. + - Raise an error if deprecated arguments are used, depending on the specified conditions. + + By default, the decorator notifies the user about the deprecated argument while the `fla.__version__` < specified `version` + in the decorator. To keep notifications with any version `warn_if_greater_or_equal_version=True` can be set. + + Args: + old_name (`str`): + Name of the deprecated keyword argument. + version (`str`): + The version in which the keyword argument was (or will be) deprecated. + new_name (`Optional[str]`, *optional*): + The new name for the deprecated keyword argument. + If specified, the deprecated keyword argument will be replaced with this new name. + warn_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): + Whether to show warning if current `fla` version is greater or equal to the deprecated version. + raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): + Whether to raise `ValueError` if current `fla` version is greater or equal to the deprecated version. + raise_if_both_names (`bool`, *optional*, defaults to `False`): + Whether to raise `ValueError` if both deprecated and new keyword arguments are set. + additional_message (`Optional[str]`, *optional*): + An additional message to append to the default deprecation message. + + Raises: + ValueError: + If `raise_if_greater_or_equal_version` is `True` and the current version >= the deprecated one, + or if `raise_if_both_names` is `True` and both old and new keyword arguments are provided. + + Returns: + Callable: + A wrapped function that handles the deprecated keyword arguments according to the specified parameters. + + Example usage with renaming argument: + + ```python + @deprecate_kwarg("reduce_labels", new_name="do_reduce_labels", version="6.0.0") + def my_function(do_reduce_labels): + print(do_reduce_labels) + + my_function(reduce_labels=True) # Will show a deprecation warning and use do_reduce_labels=True + ``` + + Example usage without renaming argument: + + ```python + @deprecate_kwarg("max_size", version="6.0.0") + def my_function(max_size): + print(max_size) + + my_function(max_size=1333) # Will show a deprecation warning + ``` + + """ + deprecated_version = version.parse(version) + current_version = version.parse(__version__) + is_greater_or_equal_version = current_version >= deprecated_version + + if is_greater_or_equal_version: + version_message = f"and removed starting from version {version}" + else: + version_message = f"and will be removed in version {version}" + + def wrapper(func): + # Required for better warning message + sig = inspect.signature(func) + function_named_args = set(sig.parameters.keys()) + is_instance_method = "self" in function_named_args + is_class_method = "cls" in function_named_args + + @functools.wraps(func) + def wrapped_func(*args, **kwargs): + # Get class + function name (just for better warning message) + func_name = func.__name__ + if is_instance_method: + func_name = f"{args[0].__class__.__name__}.{func_name}" + elif is_class_method: + func_name = f"{args[0].__name__}.{func_name}" + + minimum_action = Action.NONE + message = None + + # deprecated kwarg and its new version are set for function call -> replace it with new name + if old_name in kwargs and new_name in kwargs: + minimum_action = Action.RAISE if raise_if_both_names else Action.NOTIFY_ALWAYS + message = ( + f"Both `{old_name}` and `{new_name}` are set for `{func_name}`. " + f"Using `{new_name}={kwargs[new_name]}` and ignoring deprecated `{old_name}={kwargs[old_name]}`." + ) + kwargs.pop(old_name) + + # only deprecated kwarg is set for function call -> replace it with new name + elif old_name in kwargs and new_name is not None and new_name not in kwargs: + minimum_action = Action.NOTIFY + message = ( + f"`{old_name}` is deprecated {version_message} for `{func_name}`. " + f"Use `{new_name}` instead." + ) + kwargs[new_name] = kwargs.pop(old_name) + + # deprecated kwarg is not set for function call and new name is not specified -> just notify + elif old_name in kwargs: + minimum_action = Action.NOTIFY + message = f"`{old_name}` is deprecated {version_message} for `{func_name}`." + + if message is not None and additional_message is not None: + message = f"{message} {additional_message}" + + # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version) + if is_greater_or_equal_version: + # change to (NOTIFY, NOTIFY_ALWAYS) -> RAISE if specified + # in case we want to raise error for already deprecated arguments + if raise_if_greater_or_equal_version and minimum_action != Action.NONE: + minimum_action = Action.RAISE + + # change to NOTIFY -> NONE if specified (NOTIFY_ALWAYS can't be changed to NONE) + # in case we want to ignore notifications for already deprecated arguments + elif not warn_if_greater_or_equal_version and minimum_action == Action.NOTIFY: + minimum_action = Action.NONE + + # raise error or notify user + if minimum_action == Action.RAISE: + raise ValueError(message) + elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS): + # DeprecationWarning is ignored by default, so we use FutureWarning instead + warnings.warn(message, FutureWarning, stacklevel=2) + + return func(*args, **kwargs) + + return wrapped_func + + return wrapper + + +def checkpoint(fn): + def wrapper(*args, **kwargs): + return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs) + return wrapper + + +@functools.cache +def check_pytorch_version(version_s: str = '2.4') -> bool: + return version.parse(torch.__version__) >= version.parse(version_s) + + +def _cpu_device_warning(): + warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=1) + + +@functools.cache +def get_multiprocessor_count(tensor_idx: int = 0) -> int: + try: + return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['multiprocessor_count'] + except BaseException: + # Maybe we use a NPU device. + if triton.runtime.driver.active.get_current_target().backend == 'npu': + return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] + else: + return 1 + + +@functools.cache +def get_available_device() -> str: + try: + return triton.runtime.driver.active.get_current_target().backend + except BaseException: + _cpu_device_warning() + return 'cpu' + + +def map_triton_backend_to_torch_device() -> str: + backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ... + return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend) + + +# Avoid CUDA/Triton driver probing at import time. Runtime kernels still launch +# on the tensors' devices, but importing FLA should not initialize CUDA. +device_platform = os.environ.get("FLA_DEVICE_PLATFORM", "cuda") +device_name = "cuda" if device_platform in {"cuda", "hip"} else device_platform +device = "cuda" if device_platform in {"cuda", "hip"} else device_platform +device_torch_lib = getattr(torch, device, torch.cuda) + +IS_AMD = (device_platform == 'hip') +IS_INTEL = (device_platform == 'xpu') +IS_NVIDIA = (device_platform == 'cuda') +IS_INTEL_ALCHEMIST = False +IS_NVIDIA_HOPPER = IS_NVIDIA +IS_NVIDIA_BLACKWELL = False +USE_CUDA_GRAPH = False + +IS_TF32_SUPPORTED = IS_NVIDIA +IS_GATHER_SUPPORTED = hasattr(triton.language, 'gather') +IS_TMA_SUPPORTED = False + +if IS_NVIDIA and not IS_TF32_SUPPORTED: + # Make old card happy, since triton will use tf32 by default. + # This is a workaround for old nvidia card. + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + +if IS_TMA_SUPPORTED: + logger.info('TMA is supported, using TMA by default.') + + def alloc_fn(size: int, alignment: int, stream: int | None): + return torch.empty(size, device=torch.device(device_name, device_torch_lib.current_device()), dtype=torch.int8) + + triton.set_allocator(alloc_fn) + + +def get_all_max_shared_mem(): + try: + return [ + triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem'] + for i in range(device_torch_lib.device_count()) + ] + except BaseException: + _cpu_device_warning() + return [-1] + + +class Backend(Enum): + ADA = 101376 # RTX 4090 + AMPERE = 166912 # A100 + HOPPER = 232448 # H100 + DEFAULT = 102400 # Default + + @classmethod + def get_shared_memory(cls, arch: str) -> int: + try: + return cls[arch.upper()].value + except KeyError: + return cls.DEFAULT.value + + +@functools.cache +def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: + try: + device_shared_mem_list = get_all_max_shared_mem() + max_shared_memory = device_shared_mem_list[tensor_idx] + return max_shared_memory >= Backend.get_shared_memory(arch) + except Exception: + return False + + +if check_pytorch_version('2.4'): + device = 'cuda' if device == 'cpu' else device + autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) + autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) + + def custom_device_ctx(index: int): + return device_torch_lib.device(index) +else: + assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.' + autocast_custom_fwd = device_torch_lib.amp.custom_fwd + autocast_custom_bwd = device_torch_lib.amp.custom_bwd + + def custom_device_ctx(index: int): + return torch.cuda.device(index) + + +def _register_aliases(): + current_module = sys.modules[__name__] + for key in ( + 'IS_AMD', + 'IS_INTEL', + 'IS_NVIDIA', + 'IS_INTEL_ALCHEMIST', + 'IS_NVIDIA_HOPPER', + 'IS_NVIDIA_BLACKWELL', + 'USE_CUDA_GRAPH', + 'IS_TF32_SUPPORTED', + 'IS_GATHER_SUPPORTED', + 'IS_TMA_SUPPORTED', + ): + if hasattr(current_module, key): + setattr(current_module, key.lower(), getattr(current_module, key)) + + +_register_aliases() + +del _register_aliases diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000000000000000000000000000000000000..953320f6c78f5d2e56c2dd683ca9322be0d6cf38 --- /dev/null +++ b/generation_config.json @@ -0,0 +1,9 @@ +{ + "_from_model_config": true, + "eos_token_id": 156892, + "output_attentions": false, + "output_hidden_states": false, + "pad_token_id": 156892, + "transformers_version": "5.10.2", + "use_cache": false +} diff --git a/model-00001-of-00020.safetensors b/model-00001-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..54b87073ab7d29c48349f3433834edaa8493a09d --- /dev/null +++ b/model-00001-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea35cabb6bcb9a3d98af7c2d23b74a6c1c16d7b2250d1b33f4adf1148e1527cc +size 3644870344 diff --git a/model-00002-of-00020.safetensors b/model-00002-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7f294959e7b62f3ab7b26927af201b2e5750f706 --- /dev/null +++ b/model-00002-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d46fc3c261789b35418e83d44ddb7e14b57d90f5357e9245285fdc586cc02735 +size 3277868552 diff --git a/model-00003-of-00020.safetensors b/model-00003-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..70140ec1170c2438a1b3768ce208aceab3e14f05 --- /dev/null +++ b/model-00003-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c591c0a0e142188c63be1c33dc718a3e9ee3c3ba64ce8756feaaa6559d35cad9 +size 3277868552 diff --git a/model-00004-of-00020.safetensors b/model-00004-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..258138ad6dfcb950471c9bd5cca03ab41c6e3278 --- /dev/null +++ b/model-00004-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c72452db4fa323c790cc732f731d0a5eecbbef21bcf333083e0c08fbb7059a35 +size 3364272708 diff --git a/model-00005-of-00020.safetensors b/model-00005-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8a6f4b9bdb90cbafdd8d27a6227db466c6d13449 --- /dev/null +++ b/model-00005-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7483399d4af1b189702dfe1c54f6d93bc9eb0544acbf0018f5c631176ece89a4 +size 3328613356 diff --git a/model-00006-of-00020.safetensors b/model-00006-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c120fd620e27dfb0f45d906cbd5d7f38bf3198e7 --- /dev/null +++ b/model-00006-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a5153b37b6f06b6e2fdceff1e115513e1dbe713f824c9bca73ddaf268971130 +size 3364272708 diff --git a/model-00007-of-00020.safetensors b/model-00007-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9d7db35ec1a348a9775fb9a74ce99a013915a615 --- /dev/null +++ b/model-00007-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc7f64f0ce413e15e6f1bab1b9507305a2080d754f0450017993b4924ddbeb3d +size 3364272708 diff --git a/model-00008-of-00020.safetensors b/model-00008-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c740282c96cadb3e9b3dd6320d2da4e9c96c5037 --- /dev/null +++ b/model-00008-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2116d3232611e812a81c7247e6262e38fccc129ff4be7cfd7c4c753e359eca15 +size 3336870324 diff --git a/model-00009-of-00020.safetensors b/model-00009-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5d5f496b71763151055c13ab93d7585b27eab9f1 --- /dev/null +++ b/model-00009-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3d54aa7600d0880a26c8d887448515ab70952d0bb653c75191934467347aaa34 +size 3364272708 diff --git a/model-00010-of-00020.safetensors b/model-00010-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e62d7359195316358ea89898934357255f4e4c5d --- /dev/null +++ b/model-00010-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7afc84fcfc4308a29fd9a9d4e702505c390a4bf34849ca8d10e64cba252e0736 +size 3328613244 diff --git a/model-00011-of-00020.safetensors b/model-00011-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..696c1b1d46b40eec84458dc2faadde2b1f8f13c5 --- /dev/null +++ b/model-00011-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8232672b670795f66f6b874ce344a5ef67f42b4ee5da0507bf746dd3330e3c6f +size 3364272740 diff --git a/model-00012-of-00020.safetensors b/model-00012-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..70eda31e6b40424e61356530c1093a090c4fc4b5 --- /dev/null +++ b/model-00012-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a86044f975598c5b4eb60f5208c6224fdd024b7a4e740ddb15ed9160166d0471 +size 3364272740 diff --git a/model-00013-of-00020.safetensors b/model-00013-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..04f00b3467a170b6f65a1d27aa096bd4d4fb4d8b --- /dev/null +++ b/model-00013-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3fc3f688bf2d9fc9ccca11642231d4053a8eaca1772107dd539beb41a5ea4d52 +size 3336870348 diff --git a/model-00014-of-00020.safetensors b/model-00014-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..72e3ac0384eaeabf515e7d6bad4affef879fa8f6 --- /dev/null +++ b/model-00014-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0754df10e787162c5a54d7245f63a82df6d7a26738cdb4eb7ee375693d1603a2 +size 3364272740 diff --git a/model-00015-of-00020.safetensors b/model-00015-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..66f50f88e31aa27a64045e9b4cdef39a2091a5e3 --- /dev/null +++ b/model-00015-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0bbb28aadeab35e80ad577864c680a90574132b1f2a76803cce54887457c0123 +size 3328613388 diff --git a/model-00016-of-00020.safetensors b/model-00016-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3fdd7ad4dba4ee706b3283e7500ed1f38bac8a84 --- /dev/null +++ b/model-00016-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:edd656d1e7c6daea6e1dcbc1ef36de0a8dbb0e54fa39211db639ca02ce5e7220 +size 3364272740 diff --git a/model-00017-of-00020.safetensors b/model-00017-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0f2d459723c2ecc3578eed51d0a148d3dc3c6c91 --- /dev/null +++ b/model-00017-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:239998760c692b641107311e31b70ca42f89777973c390f71c5ed20f96bd01df +size 3364272740 diff --git a/model-00018-of-00020.safetensors b/model-00018-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..449f77aaa8f27f5762e62fc76849951b3adca5fa --- /dev/null +++ b/model-00018-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d948b2ec33dc3166c5af404c67d5aee89a1ad4b3242463121a91cb92083d5f48 +size 3336870348 diff --git a/model-00019-of-00020.safetensors b/model-00019-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e93a29ae1997dc85816835d5dad5bc703a49ea45 --- /dev/null +++ b/model-00019-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e8c98b8722eafa5ac5f957ee1f06163c745aff1cd3963f0c2b91a31c1658827 +size 3364272740 diff --git a/model-00020-of-00020.safetensors b/model-00020-of-00020.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8f80f0da1a8a3ca3973f62f35230a791eaffe8b6 --- /dev/null +++ b/model-00020-of-00020.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:909c38d3a8b436f83a388ef94089330645ce82b2eb8b9f84d915b6831d4a4401 +size 2376100200 diff --git a/model.safetensors.index.json b/model.safetensors.index.json new file mode 100644 index 0000000000000000000000000000000000000000..ffff69ae3f8660c582885f87b3b99f736b233f29 --- /dev/null +++ b/model.safetensors.index.json @@ -0,0 +1,569 @@ +{ + "metadata": { + "total_size": 66215816960 + }, + "weight_map": { + "lm_head.weight": "model-00020-of-00020.safetensors", + "model.layers.0.attention.dense.weight": "model-00001-of-00020.safetensors", + "model.layers.0.attention.key_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.0.attention.query_key_value.weight": "model-00001-of-00020.safetensors", + "model.layers.0.attention.query_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.0.input_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.0.mlp.down_proj.weight": "model-00001-of-00020.safetensors", + "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00020.safetensors", + "model.layers.0.mlp.up_proj.weight": "model-00001-of-00020.safetensors", + "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.1.attention.dense.weight": "model-00001-of-00020.safetensors", + "model.layers.1.attention.key_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.1.attention.query_key_value.weight": "model-00001-of-00020.safetensors", + "model.layers.1.attention.query_layernorm.weight": "model-00001-of-00020.safetensors", + "model.layers.1.input_layernorm.weight": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.experts_w12": "model-00001-of-00020.safetensors", + "model.layers.1.mlp.experts_w3": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.gate.expert_bias": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.gate.weight": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.shared_experts.down_proj.weight": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.shared_experts.gate_proj.weight": "model-00002-of-00020.safetensors", + "model.layers.1.mlp.shared_experts.up_proj.weight": "model-00002-of-00020.safetensors", + "model.layers.1.post_attention_layernorm.weight": "model-00002-of-00020.safetensors", + "model.layers.10.attention.branch_global_output_gain": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_local_window_mix_logit": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_mix_logits": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_output_adapter_down.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_output_adapter_up.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_output_channel_gain": "model-00010-of-00020.safetensors", + "model.layers.10.attention.branch_output_gain": "model-00010-of-00020.safetensors", + "model.layers.10.attention.dense.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.key_layernorm.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.query_key_value.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.query_layernorm.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.A_log": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.a_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.dt_bias": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.g_norm.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.k_norm.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.k_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.o_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.q_norm.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.q_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.r_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.raven_attention.v_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.10.attention.replace_alpha_raw": "model-00010-of-00020.safetensors", + "model.layers.10.input_layernorm.weight": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.experts_w12": "model-00010-of-00020.safetensors", + "model.layers.10.mlp.experts_w3": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.gate.expert_bias": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.gate.weight": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.shared_experts.down_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.shared_experts.gate_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.10.mlp.shared_experts.up_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.10.post_attention_layernorm.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_global_output_gain": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_local_window_mix_logit": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_mix_logits": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_output_adapter_down.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_output_adapter_up.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_output_channel_gain": "model-00011-of-00020.safetensors", + "model.layers.11.attention.branch_output_gain": "model-00011-of-00020.safetensors", + "model.layers.11.attention.dense.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.key_layernorm.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.A_log": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.b_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.dt_bias": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.f_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.g_proj.0.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.g_proj.1.bias": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.g_proj.1.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.k_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.o_norm.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.o_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.q_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.quasar_attention.v_proj.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.query_key_value.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.query_layernorm.weight": "model-00011-of-00020.safetensors", + "model.layers.11.attention.replace_alpha_raw": "model-00011-of-00020.safetensors", + "model.layers.11.input_layernorm.weight": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.experts_w12": "model-00011-of-00020.safetensors", + "model.layers.11.mlp.experts_w3": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.gate.expert_bias": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.gate.weight": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.shared_experts.down_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.shared_experts.gate_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.11.mlp.shared_experts.up_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.11.post_attention_layernorm.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_global_output_gain": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_local_window_mix_logit": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_mix_logits": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_output_adapter_down.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_output_adapter_up.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_output_channel_gain": "model-00012-of-00020.safetensors", + "model.layers.12.attention.branch_output_gain": "model-00012-of-00020.safetensors", + "model.layers.12.attention.dense.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.key_layernorm.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.A_log": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.b_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.dt_bias": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.f_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.g_proj.0.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.g_proj.1.bias": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.g_proj.1.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.k_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.o_norm.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.o_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.q_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.quasar_attention.v_proj.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.query_key_value.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.query_layernorm.weight": "model-00012-of-00020.safetensors", + "model.layers.12.attention.replace_alpha_raw": "model-00012-of-00020.safetensors", + "model.layers.12.input_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.experts_w12": "model-00012-of-00020.safetensors", + "model.layers.12.mlp.experts_w3": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.gate.expert_bias": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.gate.weight": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.shared_experts.down_proj.weight": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.shared_experts.gate_proj.weight": "model-00013-of-00020.safetensors", + "model.layers.12.mlp.shared_experts.up_proj.weight": "model-00013-of-00020.safetensors", + "model.layers.12.post_attention_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_global_output_gain": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_local_window_mix_logit": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_mix_logits": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_output_adapter_down.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_output_adapter_up.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_output_channel_gain": "model-00013-of-00020.safetensors", + "model.layers.13.attention.branch_output_gain": "model-00013-of-00020.safetensors", + "model.layers.13.attention.dense.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.dense.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.g_norm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.g_proj.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.key_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.query_key_value.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.query_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.gla_attention.slope": "model-00013-of-00020.safetensors", + "model.layers.13.attention.key_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.query_key_value.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.query_layernorm.weight": "model-00013-of-00020.safetensors", + "model.layers.13.attention.replace_alpha_raw": "model-00013-of-00020.safetensors", + "model.layers.13.input_layernorm.weight": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.experts_w12": "model-00013-of-00020.safetensors", + "model.layers.13.mlp.experts_w3": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.gate.expert_bias": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.gate.weight": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.shared_experts.down_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.shared_experts.gate_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.13.mlp.shared_experts.up_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.13.post_attention_layernorm.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_global_output_gain": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_local_window_mix_logit": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_mix_logits": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_output_adapter_down.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_output_adapter_up.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_output_channel_gain": "model-00014-of-00020.safetensors", + "model.layers.14.attention.branch_output_gain": "model-00014-of-00020.safetensors", + "model.layers.14.attention.dense.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.key_layernorm.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.A_log": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.b_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.dt_bias": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.f_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.g_proj.0.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.g_proj.1.bias": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.g_proj.1.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.k_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.o_norm.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.o_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.q_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.quasar_attention.v_proj.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.query_key_value.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.query_layernorm.weight": "model-00014-of-00020.safetensors", + "model.layers.14.attention.replace_alpha_raw": "model-00014-of-00020.safetensors", + "model.layers.14.input_layernorm.weight": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.experts_w12": "model-00014-of-00020.safetensors", + "model.layers.14.mlp.experts_w3": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.gate.expert_bias": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.gate.weight": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.shared_experts.down_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.shared_experts.gate_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.14.mlp.shared_experts.up_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.14.post_attention_layernorm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_global_output_gain": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_local_window_mix_logit": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_mix_logits": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_output_adapter_down.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_output_adapter_up.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_output_channel_gain": "model-00015-of-00020.safetensors", + "model.layers.15.attention.branch_output_gain": "model-00015-of-00020.safetensors", + "model.layers.15.attention.dense.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.key_layernorm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.query_key_value.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.query_layernorm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.A_log": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.a_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.dt_bias": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.g_norm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.k_norm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.k_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.o_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.q_norm.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.q_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.r_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.raven_attention.v_proj.weight": "model-00015-of-00020.safetensors", + "model.layers.15.attention.replace_alpha_raw": "model-00015-of-00020.safetensors", + "model.layers.15.input_layernorm.weight": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.experts_w12": "model-00015-of-00020.safetensors", + "model.layers.15.mlp.experts_w3": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.gate.expert_bias": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.gate.weight": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.shared_experts.down_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.shared_experts.gate_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.15.mlp.shared_experts.up_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.15.post_attention_layernorm.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_global_output_gain": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_local_window_mix_logit": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_mix_logits": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_output_adapter_down.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_output_adapter_up.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_output_channel_gain": "model-00016-of-00020.safetensors", + "model.layers.16.attention.branch_output_gain": "model-00016-of-00020.safetensors", + "model.layers.16.attention.dense.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.key_layernorm.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.A_log": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.b_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.dt_bias": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.f_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.g_proj.0.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.g_proj.1.bias": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.g_proj.1.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.k_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.o_norm.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.o_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.q_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.quasar_attention.v_proj.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.query_key_value.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.query_layernorm.weight": "model-00016-of-00020.safetensors", + "model.layers.16.attention.replace_alpha_raw": "model-00016-of-00020.safetensors", + "model.layers.16.input_layernorm.weight": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.experts_w12": "model-00016-of-00020.safetensors", + "model.layers.16.mlp.experts_w3": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.gate.expert_bias": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.gate.weight": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.shared_experts.down_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.shared_experts.gate_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.16.mlp.shared_experts.up_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.16.post_attention_layernorm.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_global_output_gain": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_local_window_mix_logit": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_mix_logits": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_output_adapter_down.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_output_adapter_up.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_output_channel_gain": "model-00017-of-00020.safetensors", + "model.layers.17.attention.branch_output_gain": "model-00017-of-00020.safetensors", + "model.layers.17.attention.dense.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.key_layernorm.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.A_log": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.b_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.dt_bias": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.f_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.g_proj.0.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.g_proj.1.bias": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.g_proj.1.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.k_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.o_norm.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.o_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.q_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.quasar_attention.v_proj.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.query_key_value.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.query_layernorm.weight": "model-00017-of-00020.safetensors", + "model.layers.17.attention.replace_alpha_raw": "model-00017-of-00020.safetensors", + "model.layers.17.input_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.experts_w12": "model-00017-of-00020.safetensors", + "model.layers.17.mlp.experts_w3": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.gate.expert_bias": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.gate.weight": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.shared_experts.down_proj.weight": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.shared_experts.gate_proj.weight": "model-00018-of-00020.safetensors", + "model.layers.17.mlp.shared_experts.up_proj.weight": "model-00018-of-00020.safetensors", + "model.layers.17.post_attention_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_global_output_gain": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_local_window_mix_logit": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_mix_logits": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_output_adapter_down.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_output_adapter_up.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_output_channel_gain": "model-00018-of-00020.safetensors", + "model.layers.18.attention.branch_output_gain": "model-00018-of-00020.safetensors", + "model.layers.18.attention.dense.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.dense.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.g_norm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.g_proj.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.key_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.query_key_value.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.query_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.gla_attention.slope": "model-00018-of-00020.safetensors", + "model.layers.18.attention.key_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.query_key_value.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.query_layernorm.weight": "model-00018-of-00020.safetensors", + "model.layers.18.attention.replace_alpha_raw": "model-00018-of-00020.safetensors", + "model.layers.18.input_layernorm.weight": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.experts_w12": "model-00018-of-00020.safetensors", + "model.layers.18.mlp.experts_w3": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.gate.expert_bias": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.gate.weight": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.shared_experts.down_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.shared_experts.gate_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.18.mlp.shared_experts.up_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.18.post_attention_layernorm.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_global_output_gain": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_local_window_mix_logit": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_mix_logits": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_output_adapter_down.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_output_adapter_up.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_output_channel_gain": "model-00019-of-00020.safetensors", + "model.layers.19.attention.branch_output_gain": "model-00019-of-00020.safetensors", + "model.layers.19.attention.dense.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.key_layernorm.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.A_log": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.b_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.dt_bias": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.f_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.g_proj.0.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.g_proj.1.bias": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.g_proj.1.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.k_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.o_norm.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.o_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.q_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.quasar_attention.v_proj.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.query_key_value.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.query_layernorm.weight": "model-00019-of-00020.safetensors", + "model.layers.19.attention.replace_alpha_raw": "model-00019-of-00020.safetensors", + "model.layers.19.input_layernorm.weight": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.experts_w12": "model-00019-of-00020.safetensors", + "model.layers.19.mlp.experts_w3": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.gate.expert_bias": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.gate.weight": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.shared_experts.down_proj.weight": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.shared_experts.gate_proj.weight": "model-00020-of-00020.safetensors", + "model.layers.19.mlp.shared_experts.up_proj.weight": "model-00020-of-00020.safetensors", + "model.layers.19.post_attention_layernorm.weight": "model-00020-of-00020.safetensors", + "model.layers.2.attention.dense.weight": "model-00002-of-00020.safetensors", + "model.layers.2.attention.key_layernorm.weight": "model-00002-of-00020.safetensors", + "model.layers.2.attention.query_key_value.weight": "model-00002-of-00020.safetensors", + "model.layers.2.attention.query_layernorm.weight": "model-00002-of-00020.safetensors", + "model.layers.2.input_layernorm.weight": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.experts_w12": "model-00002-of-00020.safetensors", + "model.layers.2.mlp.experts_w3": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.gate.expert_bias": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.gate.weight": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.shared_experts.down_proj.weight": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.shared_experts.gate_proj.weight": "model-00003-of-00020.safetensors", + "model.layers.2.mlp.shared_experts.up_proj.weight": "model-00003-of-00020.safetensors", + "model.layers.2.post_attention_layernorm.weight": "model-00003-of-00020.safetensors", + "model.layers.3.attention.dense.weight": "model-00003-of-00020.safetensors", + "model.layers.3.attention.key_layernorm.weight": "model-00003-of-00020.safetensors", + "model.layers.3.attention.query_key_value.weight": "model-00003-of-00020.safetensors", + "model.layers.3.attention.query_layernorm.weight": "model-00003-of-00020.safetensors", + "model.layers.3.input_layernorm.weight": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.experts_w12": "model-00003-of-00020.safetensors", + "model.layers.3.mlp.experts_w3": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.gate.expert_bias": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.gate.weight": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.shared_experts.down_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.shared_experts.gate_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.3.mlp.shared_experts.up_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.3.post_attention_layernorm.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_global_output_gain": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_local_window_mix_logit": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_mix_logits": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_output_adapter_down.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_output_adapter_up.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_output_channel_gain": "model-00004-of-00020.safetensors", + "model.layers.4.attention.branch_output_gain": "model-00004-of-00020.safetensors", + "model.layers.4.attention.dense.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.key_layernorm.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.A_log": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.b_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.dt_bias": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.f_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.g_proj.0.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.g_proj.1.bias": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.g_proj.1.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.k_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.o_norm.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.o_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.q_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.quasar_attention.v_proj.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.query_key_value.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.query_layernorm.weight": "model-00004-of-00020.safetensors", + "model.layers.4.attention.replace_alpha_raw": "model-00004-of-00020.safetensors", + "model.layers.4.input_layernorm.weight": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.experts_w12": "model-00004-of-00020.safetensors", + "model.layers.4.mlp.experts_w3": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.gate.expert_bias": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.gate.weight": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.shared_experts.down_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.shared_experts.gate_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.4.mlp.shared_experts.up_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.4.post_attention_layernorm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_global_output_gain": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_local_window_mix_logit": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_mix_logits": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_output_adapter_down.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_output_adapter_up.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_output_channel_gain": "model-00005-of-00020.safetensors", + "model.layers.5.attention.branch_output_gain": "model-00005-of-00020.safetensors", + "model.layers.5.attention.dense.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.key_layernorm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.query_key_value.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.query_layernorm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.A_log": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.a_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.dt_bias": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.g_norm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.k_norm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.k_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.o_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.q_norm.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.q_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.r_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.raven_attention.v_proj.weight": "model-00005-of-00020.safetensors", + "model.layers.5.attention.replace_alpha_raw": "model-00005-of-00020.safetensors", + "model.layers.5.input_layernorm.weight": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.experts_w12": "model-00005-of-00020.safetensors", + "model.layers.5.mlp.experts_w3": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.gate.expert_bias": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.gate.weight": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.shared_experts.down_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.shared_experts.gate_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.5.mlp.shared_experts.up_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.5.post_attention_layernorm.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_global_output_gain": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_local_window_mix_logit": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_mix_logits": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_output_adapter_down.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_output_adapter_up.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_output_channel_gain": "model-00006-of-00020.safetensors", + "model.layers.6.attention.branch_output_gain": "model-00006-of-00020.safetensors", + "model.layers.6.attention.dense.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.key_layernorm.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.A_log": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.b_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.dt_bias": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.f_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.g_proj.0.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.g_proj.1.bias": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.g_proj.1.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.k_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.o_norm.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.o_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.q_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.quasar_attention.v_proj.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.query_key_value.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.query_layernorm.weight": "model-00006-of-00020.safetensors", + "model.layers.6.attention.replace_alpha_raw": "model-00006-of-00020.safetensors", + "model.layers.6.input_layernorm.weight": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.experts_w12": "model-00006-of-00020.safetensors", + "model.layers.6.mlp.experts_w3": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.gate.expert_bias": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.gate.weight": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.shared_experts.down_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.shared_experts.gate_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.6.mlp.shared_experts.up_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.6.post_attention_layernorm.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_global_output_gain": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_local_window_mix_logit": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_mix_logits": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_output_adapter_down.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_output_adapter_up.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_output_channel_gain": "model-00007-of-00020.safetensors", + "model.layers.7.attention.branch_output_gain": "model-00007-of-00020.safetensors", + "model.layers.7.attention.dense.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.key_layernorm.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.A_log": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.b_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.dt_bias": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.f_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.g_proj.0.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.g_proj.1.bias": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.g_proj.1.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.k_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.o_norm.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.o_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.q_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.quasar_attention.v_proj.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.query_key_value.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.query_layernorm.weight": "model-00007-of-00020.safetensors", + "model.layers.7.attention.replace_alpha_raw": "model-00007-of-00020.safetensors", + "model.layers.7.input_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.experts_w12": "model-00007-of-00020.safetensors", + "model.layers.7.mlp.experts_w3": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.gate.expert_bias": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.gate.weight": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.shared_experts.down_proj.weight": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.shared_experts.gate_proj.weight": "model-00008-of-00020.safetensors", + "model.layers.7.mlp.shared_experts.up_proj.weight": "model-00008-of-00020.safetensors", + "model.layers.7.post_attention_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_global_output_gain": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_local_window_mix_logit": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_mix_logits": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_output_adapter_down.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_output_adapter_up.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_output_channel_gain": "model-00008-of-00020.safetensors", + "model.layers.8.attention.branch_output_gain": "model-00008-of-00020.safetensors", + "model.layers.8.attention.dense.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.dense.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.g_norm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.g_proj.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.key_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.query_key_value.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.query_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.gla_attention.slope": "model-00008-of-00020.safetensors", + "model.layers.8.attention.key_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.query_key_value.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.query_layernorm.weight": "model-00008-of-00020.safetensors", + "model.layers.8.attention.replace_alpha_raw": "model-00008-of-00020.safetensors", + "model.layers.8.input_layernorm.weight": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.experts_w12": "model-00008-of-00020.safetensors", + "model.layers.8.mlp.experts_w3": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.gate.expert_bias": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.gate.weight": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.shared_experts.down_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.shared_experts.gate_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.8.mlp.shared_experts.up_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.8.post_attention_layernorm.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_global_output_gain": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_local_window_mix_logit": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_mix_logits": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_output_adapter_down.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_output_adapter_up.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_output_channel_gain": "model-00009-of-00020.safetensors", + "model.layers.9.attention.branch_output_gain": "model-00009-of-00020.safetensors", + "model.layers.9.attention.dense.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.key_layernorm.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.A_log": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.b_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.dt_bias": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.f_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.g_proj.0.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.g_proj.1.bias": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.g_proj.1.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.k_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.o_norm.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.o_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.q_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.quasar_attention.v_proj.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.query_key_value.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.query_layernorm.weight": "model-00009-of-00020.safetensors", + "model.layers.9.attention.replace_alpha_raw": "model-00009-of-00020.safetensors", + "model.layers.9.input_layernorm.weight": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.experts_w12": "model-00009-of-00020.safetensors", + "model.layers.9.mlp.experts_w3": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.gate.expert_bias": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.gate.weight": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.shared_experts.down_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.shared_experts.gate_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.9.mlp.shared_experts.up_proj.weight": "model-00010-of-00020.safetensors", + "model.layers.9.post_attention_layernorm.weight": "model-00010-of-00020.safetensors", + "model.norm.weight": "model-00020-of-00020.safetensors", + "model.rotary_emb.inv_freq": "model-00020-of-00020.safetensors", + "model.word_embeddings.weight": "model-00001-of-00020.safetensors" + } +} diff --git a/modeling_quasar_long.py b/modeling_quasar_long.py new file mode 100644 index 0000000000000000000000000000000000000000..ee9551a31f610c4cff72d6b30fdeba446a088eff --- /dev/null +++ b/modeling_quasar_long.py @@ -0,0 +1,3227 @@ +# coding=utf-8 +# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Quasar Long model.""" + +import math +import os +import warnings +from contextlib import nullcontext +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +try: + from transformers.modeling_attn_mask_utils import ( + _prepare_4d_attention_mask, + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, + ) +except ImportError: + # transformers 5.x removed these helpers + def _prepare_4d_attention_mask(mask, dtype, tgt_len=None): + raise NotImplementedError("_prepare_4d_attention_mask removed in transformers 5.x") + def _prepare_4d_causal_attention_mask(*args, **kwargs): + raise NotImplementedError("_prepare_4d_causal_attention_mask removed in transformers 5.x") + def _prepare_4d_causal_attention_mask_for_sdpa(*args, **kwargs): + raise NotImplementedError("_prepare_4d_causal_attention_mask_for_sdpa removed in transformers 5.x") +from transformers.modeling_outputs import MoeModelOutputWithPast +try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +except ImportError: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + def dynamic_rope_update(fn): + return fn +from transformers.modeling_utils import PreTrainedModel +try: + from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13 +except ImportError: + ALL_LAYERNORM_LAYERS = [] + is_torch_greater_or_equal_than_1_13 = True # torch >= 1.13 is guaranteed in any modern env +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +# is_torch_fx_available was removed in transformers 5.x; define a no-op stub +try: + from transformers.utils.import_utils import is_torch_fx_available +except ImportError: + def is_torch_fx_available(): + return False +from .configuration_quasar_long import QuasarLongConfig +from transformers.generation.utils import GenerationMixin +from dataclasses import dataclass +from transformers.utils import ModelOutput +try: + from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss +except Exception: + LigerFusedLinearCrossEntropyLoss = None + +# ── Engram: conditional N-gram memory (DeepSeek-AI, arXiv:2601.07372) ───────── +try: + import sys as _sys + import os as _os + _HERE = _os.path.dirname(_os.path.abspath(__file__)) + if _HERE not in _sys.path: + _sys.path.insert(0, _HERE) + _RAVEN_PATH = _os.path.join(_HERE, "raven") + if _RAVEN_PATH not in _sys.path: + _sys.path.insert(0, _RAVEN_PATH) + from engram import EngramModule + _ENGRAM_AVAILABLE = True +except Exception as _engram_import_err: # pragma: no cover + EngramModule = None # type: ignore[assignment,misc] + _ENGRAM_AVAILABLE = False +def _debug_assert_finite(name: str, tensor: torch.Tensor, layer_idx: Optional[int] = None): + return + + +def _sanitize_hybrid_tensor(name: str, tensor: torch.Tensor, layer_idx: Optional[int] = None): + return tensor + + +def roll_tensor(tensor, shifts=-1, dims=-1, fill_value=0): + """Roll the tensor input along the given dimension(s). + Inserted elements are set to be 0.0. + """ + rolled_tensor = torch.roll(tensor, shifts=shifts, dims=dims) + rolled_tensor.select(dims, shifts).fill_(fill_value) + return rolled_tensor, rolled_tensor.sum() + + +@dataclass +class MoEV2CausalLMOutputWithPast(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs as well as Mixture of Expert's router hidden + states terms, to train a MoE model. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + z_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + z_loss for the sparse modules. + aux_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + aux_loss for the sparse modules. + router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the encoder model, useful to compute the auxiliary loss and the z_loss for the sparse + modules. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Cache] = None + hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None + attentions: Optional[tuple[torch.FloatTensor, ...]] = None + z_loss: Optional[torch.FloatTensor] = None + aux_loss: Optional[torch.FloatTensor] = None + router_logits: Optional[tuple[torch.FloatTensor]] = None + mtp_loss: Optional[torch.FloatTensor] = None + mtp_logits: Optional[tuple[torch.FloatTensor, ...]] = None + branch_past_key_values: Optional["QGRBranchCache"] = None + branch_mimic_loss: Optional[torch.FloatTensor] = None + branch_mimic_stats: Optional[dict] = None + + +class MoeV2ModelOutputWithPast(MoeModelOutputWithPast): + + def __init__( + self, + mtp_hidden_states=None, + branch_past_key_values=None, + branch_mimic_loss=None, + branch_mimic_stats=None, + **kwargs, + ): + super().__init__(**kwargs) + self.mtp_hidden_states = mtp_hidden_states + self.branch_past_key_values = branch_past_key_values + self.branch_mimic_loss = branch_mimic_loss + self.branch_mimic_stats = branch_mimic_stats + + +class QGRBranchCache: + """Recurrent-state cache for chunked Quasar/GLA/Raven training. + + It intentionally carries only linear/recurrent branch state, not dense GQA + KV tensors. That lets a multi-million-token logical sequence be processed + as chunks without allocating a dense multi-million-token attention cache. + """ + + def __init__(self, seen_tokens: int = 0): + self.seen_tokens = int(seen_tokens) + self.layers: list[dict] = [] + self.recurrent_states: dict[int, torch.Tensor] = {} + self.conv_states: dict[int, tuple] = {} + + def __len__(self) -> int: + return len(self.layers) + + def __getitem__(self, layer_idx: int) -> dict: + return self.layers[layer_idx] + + def get_seq_length(self, layer_idx: Optional[int] = None) -> int: + return self.seen_tokens + + def update(self, layer_idx: int, recurrent_state=None, conv_state=None, offset: int = 0, **kwargs): + layer_idx = int(layer_idx) + while len(self.layers) <= layer_idx: + self.layers.append({}) + state = self.layers[layer_idx] + if recurrent_state is not None: + state["recurrent_state"] = recurrent_state + self.recurrent_states[layer_idx] = recurrent_state + if conv_state is not None: + state["conv_state"] = conv_state + self.conv_states[layer_idx] = conv_state + if offset: + self.seen_tokens += int(offset) + return self + + def detach_(self, clone: bool = False) -> "QGRBranchCache": + def _detach(value): + if torch.is_tensor(value): + value = value.detach() + return value.clone() if clone else value + if isinstance(value, tuple): + return tuple(_detach(v) for v in value) + if isinstance(value, list): + return [_detach(v) for v in value] + if isinstance(value, dict): + return {k: _detach(v) for k, v in value.items()} + return value + + self.layers = [_detach(layer) for layer in self.layers] + self.recurrent_states = {k: _detach(v) for k, v in self.recurrent_states.items()} + self.conv_states = {k: _detach(v) for k, v in self.conv_states.items()} + return self + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + warnings.warn( + "Calling `transformers.models.QuasarLong.modeling_QuasarLong._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask" + ) + return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _make_causal_mask( + input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0 +): + warnings.warn( + "Calling `transformers.models.QuasarLong.modeling_QuasarLong._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.QuasarLong.modeling_QuasarLong.AttentionMaskConverter._make_causal_mask" + ) + return AttentionMaskConverter._make_causal_mask( + input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length + ) + + +class QuasarLongRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + QuasarLongRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def reset_parameters(self) -> None: + nn.init.ones_(self.weight) + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states.to(input_dtype)).to(input_dtype) + + +class QuasarLongGroupRMSNorm(nn.Module): + def __init__(self, hidden_size, group_norm_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.group_norm_size = group_norm_size + assert hidden_size % group_norm_size == 0, "hidden_size must be divisible by group_norm_size" + self.variance_epsilon = eps + + def reset_parameters(self) -> None: + nn.init.ones_(self.weight) + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + input_shape = hidden_states.size() + group_shape = input_shape[:-1] + (self.group_norm_size, input_shape[-1] // self.group_norm_size) + hidden_states = hidden_states.view(group_shape).to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states.to(input_dtype).view(input_shape)).to(input_dtype) + + +ALL_LAYERNORM_LAYERS.append(QuasarLongRMSNorm) + + +def _quasar_long_safe_nope_enabled(config) -> bool: + return bool(getattr(config, "use_nope", False)) and getattr(config, "long_context_mode", "") == "rope_short_nope_long" + + +def _quasar_long_global_nope_enabled(config) -> bool: + return bool(getattr(config, "use_nope", False)) and not _quasar_long_safe_nope_enabled(config) + + +class QuasarLongRotaryEmbedding(nn.Module): + def __init__(self, config: QuasarLongConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + if self.rope_type in ROPE_INIT_FUNCTIONS: + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + else: + # 'default' was removed in transformers 5.x; compute standard RoPE inv_freq inline + self.rope_init_fn = None + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + rope_theta = getattr(config, "rope_theta", 10000.0) + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + self.attention_scaling = 1.0 + self.register_buffer("inv_freq", inv_freq, persistent=True) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + if _quasar_long_global_nope_enabled(self.config): + batch, seq_len = position_ids.shape + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + partial_rotary_factor = getattr(self.config, "partial_rotary_factor", 1.0) + rotary_dim = int(head_dim * partial_rotary_factor) + cos = torch.ones(batch, seq_len, rotary_dim, device=x.device, dtype=x.dtype) + sin = torch.zeros(batch, seq_len, rotary_dim, device=x.device, dtype=x.dtype) + return cos, sin + + # Auto-recover inv_freq if it contains meta-device or weight-loader garbage values + if (self.inv_freq.device != x.device or + self.inv_freq.ndim == 0 or + self.inv_freq.shape[0] == 0 or + self.inv_freq[0].item() > 2.0 or + (self.inv_freq.shape[0] > 1 and self.inv_freq[1].item() == 0.0)): + + print(f"[ROPE DEBUG] Triggered auto-recovery! Current inv_freq device: {self.inv_freq.device}, values: {self.inv_freq[:4]}", flush=True) + partial_rotary_factor = getattr(self.config, "partial_rotary_factor", 1.0) + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + rope_theta = getattr(self.config, "rope_theta", 10000.0) + self.inv_freq = (1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=x.device) / dim))).to(x.device) + print(f"[ROPE DEBUG] Recovered inv_freq: {self.inv_freq[:4]}", flush=True) + + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + cos = cos.to(dtype=x.dtype) + sin = sin.to(dtype=x.dtype) + if _quasar_long_safe_nope_enabled(self.config): + cutoff = int(getattr(self.config, "nope_after_position", 512)) + nope_mask = (position_ids >= cutoff).unsqueeze(-1) + if bool(nope_mask.any()): + cos = torch.where(nope_mask, torch.ones_like(cos), cos) + sin = torch.where(nope_mask, torch.zeros_like(sin), sin) + return cos, sin + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, 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. + 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 the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + # Keep half or full tensor for later concatenation + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + + # Apply rotary embeddings on the first half or full tensor + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + + # Concatenate back to full shape + q_embed = torch.cat([q_embed, q_pass], dim=-1) + k_embed = torch.cat([k_embed, k_pass], dim=-1) + return q_embed, k_embed + + +class QuasarLongMLP(nn.Module): + def __init__(self, config: QuasarLongConfig, intermediate_size: int): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class QuasarLongGate(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + + self.n_group = config.n_group + self.topk_group = config.topk_group + + # topk selection algorithm + self.gating_dim = config.hidden_size + self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim))) + self.routed_scaling_factor = config.routed_scaling_factor + + self.register_buffer("expert_bias", torch.zeros((self.num_experts))) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def group_limited_topk( + self, + scores: torch.Tensor, + ): + num_tokens, _ = scores.size() + # Organize the experts into groups + group_scores = scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + + # Mask the experts based on selection groups + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_tokens, self.n_group, self.num_experts // self.n_group) + .reshape(num_tokens, -1) + ) + + masked_scores = scores.masked_fill(~score_mask.bool(), float('-inf')) + probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1) + + return probs, top_indices + + def forward(self, hidden_states): + # compute gating score + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + + scores = torch.sigmoid(logits.float()).type_as(logits) + + scores_for_routing = scores + self.expert_bias + _, topk_idx = self.group_limited_topk(scores_for_routing) + + scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits) + + topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores + topk_weight = topk_weight * self.routed_scaling_factor + + return topk_idx, topk_weight.type_as(hidden_states), logits + + +class QuasarLongSparseMoeBlock(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config: QuasarLongConfig, layer_idx: int = -1): + super().__init__() + self.layer_idx = layer_idx + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + self._setup_experts() + self.gate = QuasarLongGate(config) + if config.num_shared_experts is not None: + self.shared_experts = QuasarLongMLP( + config=config, intermediate_size=config.moe_intermediate_size * config.num_shared_experts + ) + + def reset_parameters(self) -> None: + for module in self.children(): + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + + def _setup_experts(self): + self.experts_w12 = nn.Parameter(torch.zeros(self.config.num_experts, self.config.hidden_size, 2 * self.config.moe_intermediate_size)) + self.experts_w3 = nn.Parameter(torch.zeros(self.config.num_experts, self.config.moe_intermediate_size, self.config.hidden_size)) + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + w12_key = prefix + 'experts_w12' + w3_key = prefix + 'experts_w3' + + # Initialize progressive accumulation buffers on first shard arrival + if not hasattr(self, '_temp_gate_weights'): + self._temp_gate_weights = {} + self._temp_up_weights = {} + self._temp_down_weights = {} + + num_experts = self.config.num_experts + + # Intercept and pop any separate expert weights from the active state dict shard + for k in list(state_dict.keys()): + if k.startswith(prefix + 'experts.'): + parts = k[len(prefix + 'experts.'):].split('.') + expert_idx = int(parts[0]) + proj_name = parts[1] + + weight = state_dict.pop(k) + + if proj_name == 'gate_proj': + self._temp_gate_weights[expert_idx] = weight.t() + elif proj_name == 'up_proj': + self._temp_up_weights[expert_idx] = weight.t() + elif proj_name == 'down_proj': + self._temp_down_weights[expert_idx] = weight.t() + + # Once all shards have contributed their parameters, perform in-place fusion! + if (len(self._temp_gate_weights) == num_experts and + len(self._temp_up_weights) == num_experts and + len(self._temp_down_weights) == num_experts): + + gate_stacked = torch.stack([self._temp_gate_weights[i] for i in range(num_experts)]) + up_stacked = torch.stack([self._temp_up_weights[i] for i in range(num_experts)]) + down_stacked = torch.stack([self._temp_down_weights[i] for i in range(num_experts)]) + + self.experts_w12.data.copy_(torch.cat([gate_stacked, up_stacked], dim=-1)) + self.experts_w3.data.copy_(down_stacked) + + # Deallocate temporary buffers to free CPU memory + del self._temp_gate_weights + del self._temp_up_weights + del self._temp_down_weights + + # Satisfy strict loading checks by injecting the fused tensors if HF expects them + if w12_key not in state_dict: + state_dict[w12_key] = self.experts_w12.data + state_dict[w3_key] = self.experts_w3.data + + super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + + def forward(self, hidden_states): + identity = hidden_states + bsz, seq_len, h = hidden_states.shape + topk_idx, topk_weight, router_logits = self.gate(hidden_states) + + # The old inference loop scans every expert and issues many tiny GPU ops, + # which makes one-token decode extremely slow. Keep it as an escape hatch + # for debugging, but default inference to a batched expert path. + infer_all_experts = os.environ.get("QUASAR_MOE_INFER_ALL_EXPERTS", "1") == "1" + decode_only_all_experts = os.environ.get("QUASAR_MOE_INFER_ALL_EXPERTS_DECODE_ONLY", "0") == "1" + if (not self.training) and os.environ.get("QUASAR_MOE_INFER_LOOP", "0") == "1": + y = self.moe_loop(hidden_states, topk_idx, topk_weight) + elif not self.training and infer_all_experts and (not decode_only_all_experts or seq_len == 1): + y = self.moe_all_experts(hidden_states, topk_idx, topk_weight) + else: + y = self.moe_vectorized(hidden_states, topk_idx, topk_weight) + + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y, (router_logits.view(bsz, seq_len, -1), topk_idx.view(bsz, seq_len, -1)) + + def moe_loop(self, x, topk_ids, topk_weight): + bsz, seq_len, h_dim = x.shape + k = topk_ids.shape[-1] + flat_x = x.view(-1, h_dim) + flat_topk_idx = topk_ids.view(-1) + + routed_out = torch.zeros_like(flat_x) + + flat_x_repeated = flat_x.repeat_interleave(k, dim=0) + flat_topk_weight = topk_weight.view(-1, 1) + + for i in range(self.config.num_experts): + assigned_mask = (flat_topk_idx == i) + if not assigned_mask.any(): + continue + + expert_inputs = flat_x_repeated[assigned_mask] + expert_weights = flat_topk_weight[assigned_mask] + + w12 = self.experts_w12[i] + w3 = self.experts_w3[i] + + h12 = expert_inputs @ w12 + h1, h2 = h12.chunk(2, dim=-1) + h = F.silu(h1) * h2 + expert_out = h @ w3 + + weighted_out = expert_out * expert_weights + + items_indices = torch.arange(bsz * seq_len * k, device=x.device)[assigned_mask] + token_indices = items_indices // k + + routed_out.index_add_(0, token_indices, weighted_out) + + return routed_out.view(bsz, seq_len, h_dim) + + def moe_all_experts(self, x, topk_ids, topk_weight): + bsz, seq_len, h_dim = x.shape + num_tokens = bsz * seq_len + flat_x = x.reshape(num_tokens, h_dim) + # GPT-OSS style inference: compute all experts as one batched GEMM and + # gather/weight only the routed experts. This trades memory for much + # fewer tiny launches and is especially faster for one-token decode. + expert_x = flat_x.unsqueeze(0).expand(self.config.num_experts, -1, -1) + h12 = torch.bmm(expert_x, self.experts_w12) + h1, h2 = h12.chunk(2, dim=-1) + h = F.silu(h1) * h2 + expert_out = torch.bmm(h, self.experts_w3).transpose(0, 1).contiguous() + routed = expert_out.gather( + 1, + topk_ids.reshape(num_tokens, -1, 1).expand(-1, -1, h_dim), + ) + routed = routed * topk_weight.reshape(num_tokens, -1, 1).to(dtype=routed.dtype) + return routed.sum(dim=1).view(bsz, seq_len, h_dim) + + def moe_vectorized(self, x, topk_ids, topk_weight): + bsz, seq_len, h_dim = x.shape + k = topk_ids.shape[-1] + flat_x = x.view(-1, h_dim) + + w12_t = self.experts_w12 + down_w_t = self.experts_w3 + + num_experts = self.config.num_experts + flat_topk_idx = topk_ids.view(-1) + tokens_per_expert = torch.bincount(flat_topk_idx, minlength=num_experts) + + # Capacity limit: max 2.0x average tokens per expert, minimum 128 + avg_tokens = (bsz * seq_len * k) // num_experts + capacity = max(128, int(2.0 * avg_tokens)) + + sorted_indices = torch.argsort(flat_topk_idx) + token_indices = torch.arange(bsz * seq_len, device=x.device).repeat_interleave(k)[sorted_indices] + + expert_starts = torch.cat([torch.tensor([0], device=x.device), tokens_per_expert[:-1].cumsum(0)]) + intra_offsets = torch.arange(bsz * seq_len * k, device=x.device) - expert_starts.repeat_interleave(tokens_per_expert) + expert_idx = flat_topk_idx[sorted_indices] + + # Apply capacity limit mask + mask = intra_offsets < capacity + sorted_indices = sorted_indices[mask] + token_indices = token_indices[mask] + expert_idx = expert_idx[mask] + intra_offsets = intra_offsets[mask] + + kept_per_expert = torch.bincount(expert_idx, minlength=num_experts) + active_experts = torch.nonzero(kept_per_expert, as_tuple=False).flatten() + active_counts = kept_per_expert[active_experts] + active_starts = torch.cat( + [active_counts.new_zeros(1), active_counts.cumsum(0)[:-1]], + dim=0, + ) + + grouped_x = flat_x[token_indices] + gating_flat = topk_weight.view(-1) + sorted_gating = gating_flat[sorted_indices].unsqueeze(1) + routed_out = torch.zeros_like(flat_x) + + # Keep the batched-GEMM path, but tile experts to cap peak activation memory. + # Two-H200 runs leave very little headroom after FSDP unshards the MoE weights. + default_tile_size = "1" if self.training else "8" + expert_tile_size = int(os.environ.get("QUASAR_MOE_TILE_SIZE", default_tile_size)) + for tile_start in range(0, active_experts.numel(), expert_tile_size): + tile_end = min(tile_start + expert_tile_size, active_experts.numel()) + tile_experts = active_experts[tile_start:tile_end] + tile_counts = active_counts[tile_start:tile_end] + tile_capacity = int(tile_counts.max().item()) + tile_data_start = int(active_starts[tile_start].item()) + tile_data_end = int((active_starts[tile_end - 1] + active_counts[tile_end - 1]).item()) + + tile_grouped_x = grouped_x[tile_data_start:tile_data_end] + tile_token_indices = token_indices[tile_data_start:tile_data_end] + tile_intra_offsets = intra_offsets[tile_data_start:tile_data_end] + tile_gating = sorted_gating[tile_data_start:tile_data_end] + + if tile_experts.numel() == 1: + # Python-int indexing returns a view. Tensor/list indexing copies the + # expert weights, which can OOM when FSDP has already unsharded them. + expert_id = int(tile_experts[0].item()) + h12 = tile_grouped_x.matmul(w12_t[expert_id]) + h1, h2 = h12.chunk(2, dim=-1) + h = F.silu(h1) * h2 + expert_out = h.matmul(down_w_t[expert_id]) + routed_out.index_add_(0, tile_token_indices, expert_out * tile_gating) + continue + + tile_w12 = w12_t[tile_experts] + tile_w3 = down_w_t[tile_experts] + + tile_expert_positions = torch.repeat_interleave( + torch.arange(tile_experts.numel(), device=x.device), + tile_counts, + ) + + padded_x = torch.zeros( + tile_experts.numel(), + tile_capacity, + h_dim, + device=x.device, + dtype=x.dtype, + ) + padded_x_flat = padded_x.view(-1, h_dim) + flat_dest_indices = tile_expert_positions * tile_capacity + tile_intra_offsets + padded_x_flat.index_put_((flat_dest_indices,), tile_grouped_x) + + h12 = torch.bmm(padded_x, tile_w12) + h1, h2 = h12.chunk(2, dim=-1) + h = F.silu(h1) * h2 + expert_out_padded = torch.bmm(h, tile_w3) + + tile_expert_out = expert_out_padded.view(-1, h_dim)[flat_dest_indices] + weighted_out = tile_expert_out * tile_gating + routed_out.index_add_(0, tile_token_indices, weighted_out) + + return routed_out.view(bsz, seq_len, h_dim) + + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + # CRITICAL: Use .tolist() instead of .cpu().numpy() to reduce sync overhead if possible + # but the real fix is the vectorized path above. + tokens_per_expert_list = tokens_per_expert.tolist() + outputs = [] + dummy_outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert_list): + expert = self.experts[i] + if num_tokens > 0: + expert_out = expert(sorted_tokens[start_idx:start_idx+num_tokens]) + outputs.append(expert_out) + start_idx += num_tokens + else: + # Force ZeRO-3 hooks to trigger by passing a 1-element dummy tensor + # Multiply by 0.0 and sum to a scalar so it can be added to the graph safely. + dummy_input = sorted_tokens[0:1] + dummy_out = expert(dummy_input) * 0.0 + dummy_outputs.append(dummy_out.sum()) + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + + # Add the dummy outputs to the graph to prevent PyTorch from skipping the backward pass + if len(dummy_outputs) > 0: + final_out = final_out + sum(dummy_outputs) + + return final_out + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int, head_first: bool = True) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). + """ + if n_rep == 1: + return hidden_states + if head_first: + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + 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) + batch, slen, num_key_value_heads, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim) + return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->QuasarLong +class QuasarLongAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: QuasarLongConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim or self.hidden_size // self.num_heads + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + self.rope_dim = int(self.head_dim * partial_rotary_factor) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + + self.query_key_value = nn.Linear( + self.hidden_size, + (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=config.use_qkv_bias, + ) + + if self.config.use_qk_norm: + self.query_layernorm = QuasarLongRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = QuasarLongRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias) + + def reset_parameters(self) -> None: + for module in self.children(): + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def 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, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + + bsz, q_len, _ = hidden_states.size() + + qkv = self.query_key_value(hidden_states) + qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) + + query_states, key_states, value_states = qkv.split( + [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if self.config.use_qk_norm: + query_states = self.query_layernorm(query_states) + key_states = self.key_layernorm(key_states) + + cos, sin = position_embeddings + if not _quasar_long_global_nope_enabled(self.config): + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + cache_kwargs = {"sin": sin, "cos": cos} + if self.layer_idx < self.config.num_hidden_layers: + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + kv_seq_len = key_states.shape[-2] + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # 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, -1) + + attn_output = self.dense(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->QuasarLong +class QuasarLongFlashAttention2(QuasarLongAttention): + """ + QuasarLong flash attention module. This module inherits from `QuasarLongAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # QuasarLongFlashAttention2 attention does not support output_attentions + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + + qkv = self.query_key_value(hidden_states) + qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) + + query_states, key_states, value_states = qkv.split( + [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if self.config.use_qk_norm: + query_states = self.query_layernorm(query_states) + key_states = self.key_layernorm(key_states) + + cos, sin = position_embeddings + if not _quasar_long_global_nope_enabled(self.config): + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None and self.layer_idx < self.config.num_hidden_layers: + cache_kwargs = {"sin": sin, "cos": cos} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently cast in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slow down training & inference so it is recommended to not cast the LayerNorms + # in fp32. (QuasarLongRMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + # Handle the case where the model is quantized + if hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + elif torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + else: + target_dtype = self.query_key_value.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = self._flash_attention_forward( + query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.dense(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + ): + """ + 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) + query_length (`int`): + The length of the query sequence in terms of tokens. This represents the number of tokens in the + `query_states` tensor along the sequence dimension. It is used to determine the effective sequence + length for attention computations. + """ + 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 QuasarLongFlashAttention2 __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 = 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, + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal + ) + + return attn_output + + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->QuasarLong +class QuasarLongSdpaAttention(QuasarLongAttention): + """ + QuasarLong attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `QuasarLongAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from QuasarLongAttention.forward + def 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, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "QuasarLongModel is using QuasarLongSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + bsz, q_len, _ = hidden_states.size() + + qkv = self.query_key_value(hidden_states) + qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) + + query_states, key_states, value_states = qkv.split( + [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if self.config.use_qk_norm: + query_states = self.query_layernorm(query_states) + key_states = self.key_layernorm(key_states) + + cos, sin = position_embeddings + if not _quasar_long_global_nope_enabled(self.config): + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None and self.layer_idx < self.config.num_hidden_layers: + cache_kwargs = {"sin": sin, "cos": cos} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if attention_mask is not None: + kv_seq_len = key_states.shape[-2] + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=self.is_causal and attention_mask is None and q_len > 1, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.dense(attn_output) + + return attn_output, None, past_key_value + + +class QuasarLongLinearAttention(nn.Module): + """Quasar-shaped GLA branch used as the trainable replacement candidate. + + This intentionally mirrors the original attention projection path: one + fused QKV projection, optional QK RMSNorm, RoPE on Q/K, GQA-style KV repeat, + and a final dense projection back to hidden size. + """ + + def __init__(self, config: QuasarLongConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim or self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.mode = getattr(config, "hybrid_gla_mode", "chunk") + + self.query_key_value = nn.Linear( + self.hidden_size, + (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=config.use_qkv_bias, + ) + if self.config.use_qk_norm: + self.query_layernorm = QuasarLongRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = QuasarLongRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias) + self.g_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.g_norm = QuasarLongGroupRMSNorm( + self.num_heads * self.head_dim, + group_norm_size=getattr(config, "hybrid_gla_group_norm_size", self.num_heads), + eps=config.rms_norm_eps, + ) + slope = -self.build_slope_tensor(self.num_heads) + if config.num_hidden_layers > 1 and layer_idx is not None: + slope = slope * (1 - max(layer_idx - 1, 0) / (config.num_hidden_layers - 1) + 1e-5) + self.register_buffer("slope", slope, persistent=True) + + from fla.ops.simple_gla.chunk import chunk_simple_gla + from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla + from fla.ops.simple_gla.naive import naive_chunk_simple_gla, naive_recurrent_simple_gla + + self.lightning_attn_ops = { + "chunk": chunk_simple_gla, + "fused_recurrent": fused_recurrent_simple_gla, + "naive_chunk": naive_chunk_simple_gla, + "naive_recurrent": naive_recurrent_simple_gla, + } + + def reset_parameters(self) -> None: + pass + + @staticmethod + def build_slope_tensor(n_attention_heads: int): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio ** i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + return torch.tensor(get_slopes(n_attention_heads), dtype=torch.float32) + + def 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, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if past_key_value is None: + # The hybrid wrapper passes the shared QGR branch cache as + # `past_key_values` to match Quasar/Raven. Accept that alias here so + # GLA can use the recurrent one-token decode kernel instead of the + # much slower chunk kernel. + past_key_value = kwargs.get("past_key_values", None) + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding)." + ) + assert not output_attentions, "GLA replacement branch does not support output_attentions=True" + + bsz, q_len, _ = hidden_states.size() + mode = self.mode + if ( + (not self.training) + and q_len == 1 + and use_cache + and past_key_value is not None + and mode in {"chunk", "fused_chunk", "naive_chunk"} + and "fused_recurrent" in self.lightning_attn_ops + ): + mode = "fused_recurrent" + qkv = self.query_key_value(hidden_states) + _debug_assert_finite("qkv_proj", qkv, self.layer_idx) + qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) + query_states, key_states, value_states = qkv.split( + [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 + ) + _debug_assert_finite("qkv_split_q", query_states, self.layer_idx) + _debug_assert_finite("qkv_split_k", key_states, self.layer_idx) + _debug_assert_finite("qkv_split_v", value_states, self.layer_idx) + + if self.config.use_qk_norm: + query_states = self.query_layernorm(query_states) + key_states = self.key_layernorm(key_states) + _debug_assert_finite("qk_norm_q", query_states, self.layer_idx) + _debug_assert_finite("qk_norm_k", key_states, self.layer_idx) + + if position_embeddings is not None: + cos, sin = position_embeddings + if not _quasar_long_global_nope_enabled(self.config): + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=2) + _debug_assert_finite("rope_q", query_states, self.layer_idx) + _debug_assert_finite("rope_k", key_states, self.layer_idx) + + if self.num_key_value_groups > 1: + key_states = repeat_kv(key_states, self.num_key_value_groups, head_first=False) + value_states = repeat_kv(value_states, self.num_key_value_groups, head_first=False) + _debug_assert_finite("repeat_k", key_states, self.layer_idx) + _debug_assert_finite("repeat_v", value_states, self.layer_idx) + + if attention_mask is not None and not bool(attention_mask.all()): + value_states = value_states * attention_mask[:, -q_len:, None, None].to(dtype=value_states.dtype) + + recurrent_state = None + if past_key_value is not None and self.layer_idx is not None: + try: + if len(past_key_value) > self.layer_idx: + last_state = past_key_value[self.layer_idx] + if isinstance(last_state, dict): + recurrent_state = last_state.get("recurrent_state", None) + except TypeError: + pass + kernel_fp32 = bool(getattr(self.config, "hybrid_gla_kernel_fp32", False)) + kernel_dtype = torch.float32 if kernel_fp32 else query_states.dtype + query_states = query_states.to(kernel_dtype) + key_states = key_states.to(kernel_dtype) + value_states = value_states.to(kernel_dtype) + decay = self.slope.to(dtype=kernel_dtype, device=hidden_states.device) + o, recurrent_state = self.lightning_attn_ops[mode]( + q=query_states, + k=key_states, + v=value_states, + g=decay[None, None, :].expand(bsz, q_len, self.num_heads), + initial_state=recurrent_state, + output_final_state=use_cache, + ) + if past_key_value is not None and use_cache: + past_key_value.update( + layer_idx=self.layer_idx, + recurrent_state=recurrent_state, + conv_state=None, + offset=q_len, + ) + _debug_assert_finite("simple_gla_output", o, self.layer_idx) + + o = o.reshape(bsz, q_len, -1) + o = self.g_norm(o) + _debug_assert_finite("g_norm", o, self.layer_idx) + o = o * torch.sigmoid(self.g_proj(hidden_states)) + _debug_assert_finite("output_gate", o, self.layer_idx) + o = self.dense(o.to(hidden_states.dtype)) + _debug_assert_finite("dense", o, self.layer_idx) + return o, None, past_key_value + + +class QuasarLongHybridReplacementSdpaAttention(QuasarLongSdpaAttention): + """SDPA attention with a gated Quasar+GLA replacement path. + + Original GQA parameters stay at the top level of this module, so pretrained + `attention.query_key_value` and `attention.dense` weights load unchanged. + """ + + def __init__(self, config: QuasarLongConfig, layer_idx: Optional[int] = None): + super().__init__(config=config, layer_idx=layer_idx) + hybrid_layers = set(getattr(config, "hybrid_attention_layers", []) or []) + self.hybrid_enabled = layer_idx in hybrid_layers + self.hybrid_replacement_mode = str(getattr(config, "hybrid_replacement_mode", "gated")).lower() + self.last_gqa_output = None + self.last_linear_output = None + self.last_quasar_output = None + self.last_raven_output = None + self.last_gla_output = None + self.last_local_window_output = None + self.last_pre_channel_output = None + self.last_global_pre_channel_output = None + if not self.hybrid_enabled: + return + + from fla.layers.quasar import QuasarAttention + if not os.path.isdir(os.path.join(_HERE, "raven")): + raise ModuleNotFoundError("Quasar requires the bundled repo-local raven/ folder for Raven hybrid layers") + from raven.layers.raven import RavenAttention + + use_short_conv = bool(getattr(config, "hybrid_use_short_conv", False)) + self.hybrid_branch_layout = str(getattr(config, "hybrid_branch_layout", "mixed") or "mixed").strip().lower() + self.hybrid_assigned_branch = "mixed" + if self.hybrid_branch_layout == "layerwise": + enabled_branches = { + "quasar": bool(getattr(config, "hybrid_quasar_enabled", True)), + "raven": bool(getattr(config, "hybrid_raven_enabled", False)), + "gla": bool(getattr(config, "hybrid_gla_enabled", True)), + } + cycle = getattr(config, "hybrid_layerwise_cycle", ["quasar", "raven", "gla"]) or ["quasar"] + cycle = [ + str(branch).strip().lower() + for branch in cycle + if str(branch).strip().lower() in enabled_branches + and enabled_branches[str(branch).strip().lower()] + ] + if not cycle: + cycle = [name for name, enabled in enabled_branches.items() if enabled] or ["quasar"] + hybrid_order = sorted(hybrid_layers) + branch_pos = hybrid_order.index(layer_idx) if layer_idx in hybrid_order else 0 + self.hybrid_assigned_branch = cycle[branch_pos % len(cycle)] + self.replace_alpha_raw = nn.Parameter( + torch.tensor([float(getattr(config, "hybrid_alpha_init", -15.0))], dtype=torch.float32) + ) + self.branch_mix_logits = nn.Parameter(torch.zeros(3, dtype=torch.float32)) + self.branch_output_gain = nn.Parameter( + torch.tensor([float(getattr(config, "hybrid_output_gain_init", 1.0))], dtype=torch.float32) + ) + self.branch_global_output_gain = nn.Parameter( + torch.tensor([float(getattr(config, "hybrid_global_output_gain_init", getattr(config, "hybrid_output_gain_init", 1.0)))], dtype=torch.float32) + ) + self.branch_output_channel_gain = nn.Parameter(torch.ones(config.hidden_size, dtype=torch.float32)) + + local_window_layers = set(getattr(config, "hybrid_local_window_layers", []) or []) + self.local_window_size = int(getattr(config, "hybrid_local_window_size", 0) or 0) + self.local_window_enabled = self.local_window_size > 0 and ( + not local_window_layers or layer_idx in local_window_layers + ) + local_window_fraction = float(getattr(config, "hybrid_local_window_fraction", 0.0) or 0.0) + local_window_fraction = min(max(local_window_fraction, 1e-6), 1.0 - 1e-6) + self.branch_local_window_mix_logit = nn.Parameter( + torch.tensor([math.log(local_window_fraction / (1.0 - local_window_fraction))], dtype=torch.float32) + ) + local_meta_layers = set(getattr(config, "hybrid_local_meta_layers", []) or []) + self.local_meta_enabled = self.local_window_enabled and ( + not local_meta_layers or layer_idx in local_meta_layers + ) + self.local_meta_tokens = int(getattr(config, "hybrid_local_meta_tokens", 0) or 0) + if not self.local_meta_enabled: + self.local_meta_tokens = 0 + if self.local_window_enabled and self.local_meta_tokens > 0: + self.local_meta_key = nn.Parameter( + torch.empty(self.num_heads, self.local_meta_tokens, self.head_dim, dtype=torch.float32) + ) + self.local_meta_value = nn.Parameter( + torch.empty(self.num_heads, self.local_meta_tokens, self.head_dim, dtype=torch.float32) + ) + self._reset_local_meta_tokens() + else: + self.local_meta_key = None + self.local_meta_value = None + self.branch_output_adapter_rank = int(getattr(config, 'hybrid_output_adapter_rank', 16) or 0) + self.branch_output_adapter_scale = float( + getattr(config, 'hybrid_output_adapter_alpha', max(self.branch_output_adapter_rank, 1)) + ) / max(self.branch_output_adapter_rank, 1) + if self.branch_output_adapter_rank > 0: + self.branch_output_adapter_down = nn.Linear( + config.hidden_size, self.branch_output_adapter_rank, bias=False + ) + self.branch_output_adapter_up = nn.Linear( + self.branch_output_adapter_rank, config.hidden_size, bias=False + ) + self.branch_output_adapter_down._skip_quasar_hf_init = True + self.branch_output_adapter_up._skip_quasar_hf_init = True + self._reset_branch_output_adapter() + else: + self.branch_output_adapter_down = None + self.branch_output_adapter_up = None + self.distill_sum = nn.Identity() + gla_layers = set(getattr(config, "hybrid_gla_layers", []) or []) + gla_enabled_here = bool(getattr(config, "hybrid_gla_enabled", True)) and ( + not gla_layers or layer_idx in gla_layers + ) + layerwise = self.hybrid_branch_layout == "layerwise" + want_quasar = bool(getattr(config, "hybrid_quasar_enabled", True)) and ( + not layerwise or self.hybrid_assigned_branch == "quasar" + ) + want_raven = bool(getattr(config, "hybrid_raven_enabled", False)) and ( + not layerwise or self.hybrid_assigned_branch == "raven" + ) + want_gla = gla_enabled_here and ( + not layerwise or self.hybrid_assigned_branch == "gla" + ) + self.gla_attention = ( + QuasarLongLinearAttention(config=config, layer_idx=layer_idx) + if want_gla + else None + ) + self.quasar_attention = ( + QuasarAttention( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_heads=config.num_attention_heads, + mode=getattr(config, "hybrid_quasar_mode", "chunk"), + use_short_conv=use_short_conv, + conv_size=4, + conv_bias=False, + norm_eps=config.rms_norm_eps, + layer_idx=layer_idx, + ) + if want_quasar + else None + ) + self.raven_attention = ( + RavenAttention( + mode=getattr(config, "hybrid_gla_mode", "fused_recurrent"), + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + num_slots=getattr(config, "hybrid_raven_slots", 64), + topk=getattr(config, "hybrid_raven_topk", 32), + decay_type=getattr(config, "hybrid_raven_decay_type", "Mamba2"), + add_gumbel_noise=bool(getattr(config, "hybrid_raven_add_gumbel_noise", False)), + norm_eps=config.rms_norm_eps, + layer_idx=layer_idx, + ) + if want_raven + else None + ) + for branch in (self.gla_attention, self.quasar_attention, self.raven_attention): + if branch is not None: + for module in branch.modules(): + module._skip_quasar_hf_init = True + + def _reset_local_meta_tokens(self) -> None: + if self.local_meta_key is None or self.local_meta_value is None: + return + std = float(getattr(self.config, "hybrid_local_meta_init_std", 0.02) or 0.02) + nn.init.normal_(self.local_meta_key, mean=0.0, std=std) + nn.init.normal_(self.local_meta_value, mean=0.0, std=std) + + def _reset_branch_output_adapter(self) -> None: + if self.branch_output_adapter_down is None or self.branch_output_adapter_up is None: + return + nn.init.kaiming_uniform_(self.branch_output_adapter_down.weight, a=math.sqrt(5)) + self.branch_output_adapter_up.weight.data.zero_() + + def _apply_branch_output_adapter(self, linear_out: torch.Tensor) -> torch.Tensor: + if self.branch_output_adapter_down is None or self.branch_output_adapter_up is None: + return linear_out + adapter_hidden = self.branch_output_adapter_down(linear_out) + adapter_out = self.branch_output_adapter_up(adapter_hidden) + return linear_out + self.branch_output_adapter_scale * adapter_out.to(dtype=linear_out.dtype) + + @staticmethod + def _to_linear_attention_mask( + attention_mask: Optional[torch.Tensor], + *, + bsz: int, + q_len: int, + device: torch.device, + ) -> Optional[torch.Tensor]: + if attention_mask is None: + return None + if attention_mask.dim() == 2: + mask = attention_mask[:, -q_len:] + return None if bool(mask.all()) else mask.to(device=device, dtype=torch.int32) + if attention_mask.dim() == 4 and attention_mask.shape[1] == 1: + mask = attention_mask[:, 0, -1, -q_len:] + mask = (mask > -1e4) + return None if bool(mask.all()) else mask.to(device=device, dtype=torch.int32) + raise ValueError(f"Unsupported linear attention mask shape: {attention_mask.shape}") + + def reset_hybrid_branch_parameters(self) -> None: + if hasattr(self, "engram") and self.engram is not None: + self.engram._init_weights() + if not self.hybrid_enabled: + return + if self.gla_attention is not None and hasattr(self.gla_attention, "slope"): + slope = -self.gla_attention.build_slope_tensor(self.gla_attention.num_heads) + if self.config.num_hidden_layers > 1 and self.layer_idx is not None: + slope = slope * (1 - max(self.layer_idx - 1, 0) / (self.config.num_hidden_layers - 1) + 1e-5) + self.gla_attention.slope.data.copy_(slope.to(device=self.gla_attention.slope.device, dtype=self.gla_attention.slope.dtype)) + for branch in (self.gla_attention, self.quasar_attention, self.raven_attention): + if branch is None: + continue + for module in branch.modules(): + if module is branch: + continue + if isinstance(module, (QuasarLongRMSNorm, QuasarLongGroupRMSNorm)): + module.weight.data.fill_(1.0) + continue + reset = getattr(module, "reset_parameters", None) + if callable(reset): + reset() + if hasattr(branch, "A_log"): + branch.A_log.data.copy_(torch.log(torch.empty_like(branch.A_log).uniform_(1, 16))) + if hasattr(branch, "dt_bias"): + branch.dt_bias.data.zero_() + self.replace_alpha_raw.data.fill_(float(getattr(self.config, "hybrid_alpha_init", -15.0))) + self.branch_mix_logits.data.zero_() + self.branch_output_gain.data.fill_(float(getattr(self.config, "hybrid_output_gain_init", 1.0))) + self.branch_global_output_gain.data.fill_( + float(getattr(self.config, "hybrid_global_output_gain_init", getattr(self.config, "hybrid_output_gain_init", 1.0))) + ) + self.branch_output_channel_gain.data.fill_(1.0) + + local_window_fraction = float(getattr(self.config, "hybrid_local_window_fraction", 0.0) or 0.0) + local_window_fraction = min(max(local_window_fraction, 1e-6), 1.0 - 1e-6) + self.branch_local_window_mix_logit.data.fill_(math.log(local_window_fraction / (1.0 - local_window_fraction))) + self._reset_branch_output_adapter() + self._reset_local_meta_tokens() + + + def _local_window_fraction(self, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + local_fraction = torch.sigmoid(self.branch_local_window_mix_logit).to(dtype=dtype, device=device) + max_fraction = float(getattr(self.config, "hybrid_local_window_max_fraction", 0.3333333) or 0.3333333) + return torch.clamp(local_fraction, min=0.0, max=max_fraction) + + def _local_window_attention_output( + self, + hidden_states: torch.Tensor, + *, + attention_mask: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + # LoLCATs-style local softmax path. This keeps only a small causal window exact, + # while the global branch remains Quasar+GLA. + bsz, q_len, _ = hidden_states.shape + qkv = self.query_key_value(hidden_states) + qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) + query_states, key_states, value_states = qkv.split( + [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if self.config.use_qk_norm: + query_states = self.query_layernorm(query_states) + key_states = self.key_layernorm(key_states) + + if position_embeddings is not None: + cos, sin = position_embeddings + if not _quasar_long_global_nope_enabled(self.config): + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + scores = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + query_pos = torch.arange(q_len, device=hidden_states.device)[:, None] + key_pos = torch.arange(q_len, device=hidden_states.device)[None, :] + window = max(int(self.local_window_size), 1) + local_mask = (key_pos <= query_pos) & (key_pos >= query_pos - window + 1) + min_value = torch.finfo(scores.dtype).min + scores = scores.masked_fill(~local_mask.view(1, 1, q_len, q_len), min_value) + + if attention_mask is not None: + if attention_mask.dim() == 2: + key_padding_mask = attention_mask[:, -q_len:].to(device=hidden_states.device).bool() + scores = scores.masked_fill(~key_padding_mask.view(bsz, 1, 1, q_len), min_value) + elif attention_mask.dim() == 4: + scores = scores + attention_mask[:, :, -q_len:, -q_len:].to(device=scores.device, dtype=scores.dtype) + else: + raise ValueError(f"Unsupported local attention mask shape: {attention_mask.shape}") + + if self.local_meta_key is not None and self.local_meta_value is not None: + meta_key = self.local_meta_key.to(device=query_states.device, dtype=query_states.dtype) + meta_value = self.local_meta_value.to(device=value_states.device, dtype=value_states.dtype) + meta_scores = torch.einsum("bhqd,hmd->bhqm", query_states, meta_key) / math.sqrt(self.head_dim) + scores = torch.cat([meta_scores, scores], dim=-1) + meta_value = meta_value.unsqueeze(0).expand(bsz, -1, -1, -1) + value_states = torch.cat([meta_value, value_states], dim=2) + + probs = nn.functional.softmax(scores, dim=-1, dtype=torch.float32).to(query_states.dtype) + probs = nn.functional.dropout(probs, p=self.attention_dropout, training=self.training) + local_out = torch.matmul(probs, value_states) + local_out = local_out.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size) + local_out = self.dense(local_out) + _debug_assert_finite("local_window_output", local_out, self.layer_idx) + return local_out + + def _linear_attention_output( + self, + hidden_states: torch.Tensor, + *, + attention_mask: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + output_attentions: bool, + branch_past_key_values: Optional[QGRBranchCache] = None, + branch_use_cache: bool = False, + ) -> torch.Tensor: + if ( + self.training + and bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + and not torch.is_grad_enabled() + ): + with torch.enable_grad(): + return self._linear_attention_output( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + output_attentions=output_attentions, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + _debug_assert_finite("linear_input_hidden_states", hidden_states, self.layer_idx) + bsz, q_len, _ = hidden_states.shape + linear_attention_mask = self._to_linear_attention_mask( + attention_mask, + bsz=bsz, + q_len=q_len, + device=hidden_states.device, + ) + outputs = [] + self.last_quasar_output = None + self.last_raven_output = None + self.last_gla_output = None + active_branches = None + if self.training and ( + bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + or bool(getattr(self.config, "hybrid_attention_collect_branch_loss", False)) + ): + active_branches = set(getattr(self.config, "branch_mimic_branches", ["quasar", "raven", "gla", "mixed"])) + eval_force_branch = None + if not self.training: + eval_force_branch = str(getattr(self.config, "hybrid_eval_force_branch", "") or "").strip().lower() + if eval_force_branch in {"quasar", "raven", "gla", "mixed"}: + active_branches = {eval_force_branch} + needs_mixed = active_branches is None or "mixed" in active_branches + # Branch-mimic distillation trains the replacement attention modules to + # match the frozen GQA teacher on fixed hidden features. Detaching here + # prevents backward from traversing the full frozen 20B base model. + branch_hidden_states = hidden_states.detach() if active_branches is not None else hidden_states + + # 1. Quasar + if self.quasar_attention is not None and (active_branches is None or "quasar" in active_branches or needs_mixed): + use_quasar_rope = bool(getattr(self.config, "hybrid_quasar_use_rope", False)) and not _quasar_long_global_nope_enabled(self.config) + cos, sin = position_embeddings if (use_quasar_rope and position_embeddings is not None) else (None, None) + if cos is not None and sin is not None: + q_head_dim = int(self.quasar_attention.head_dim) + cos = cos[..., :q_head_dim] + sin = sin[..., :q_head_dim] + if cos.dim() == 3: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q_out = self.quasar_attention( + hidden_states=branch_hidden_states, + attention_mask=linear_attention_mask, + past_key_values=branch_past_key_values, + use_cache=branch_use_cache, + output_attentions=False, + cos=cos, + sin=sin, + )[0] + self.last_quasar_output = q_out + _debug_assert_finite("quasar_output", q_out, self.layer_idx) + q_out = _sanitize_hybrid_tensor("quasar_output", q_out, self.layer_idx) + outputs.append(q_out) + else: + outputs.append(branch_hidden_states.new_zeros(branch_hidden_states.shape)) + + # 2. Raven + if self.raven_attention is not None and (active_branches is None or "raven" in active_branches or needs_mixed): + r_out = self.raven_attention( + hidden_states=branch_hidden_states, + attention_mask=linear_attention_mask, + past_key_values=branch_past_key_values, + use_cache=branch_use_cache, + output_attentions=output_attentions, + )[0] + self.last_raven_output = r_out + _debug_assert_finite("raven_output", r_out, self.layer_idx) + r_out = _sanitize_hybrid_tensor("raven_output", r_out, self.layer_idx) + outputs.append(r_out) + else: + outputs.append(branch_hidden_states.new_zeros(branch_hidden_states.shape)) + + # 3. GLA + if self.gla_attention is not None and (active_branches is None or "gla" in active_branches or needs_mixed): + g_out = self.gla_attention( + hidden_states=branch_hidden_states, + attention_mask=linear_attention_mask, + past_key_values=branch_past_key_values, + use_cache=branch_use_cache, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + )[0] + self.last_gla_output = g_out + _debug_assert_finite("gla_output", g_out, self.layer_idx) + g_out = _sanitize_hybrid_tensor("gla_output", g_out, self.layer_idx) + outputs.append(g_out) + else: + outputs.append(branch_hidden_states.new_zeros(branch_hidden_states.shape)) + + mix = torch.softmax(self.branch_mix_logits.float(), dim=0).to(dtype=hidden_states.dtype, device=hidden_states.device) + available_mask = torch.tensor( + [ + 1.0 if self.quasar_attention is not None else 0.0, + 1.0 if self.raven_attention is not None else 0.0, + 1.0 if self.gla_attention is not None else 0.0, + ], + dtype=mix.dtype, + device=mix.device, + ) + mix = mix * available_mask + if active_branches is not None and not needs_mixed: + mask = torch.tensor( + [ + 1.0 if "quasar" in active_branches else 0.0, + 1.0 if "raven" in active_branches else 0.0, + 1.0 if "gla" in active_branches else 0.0, + ], + dtype=mix.dtype, + device=mix.device, + ) + mix = mix * mask + mix = mix / torch.clamp(mix.sum(), min=1e-6) + global_out = ( + mix[0] * outputs[0].to(dtype=hidden_states.dtype) + + mix[1] * outputs[1].to(dtype=hidden_states.dtype) + + mix[2] * outputs[2].to(dtype=hidden_states.dtype) + ) + global_out = _sanitize_hybrid_tensor("global_branch_mix", global_out, self.layer_idx) + self._last_global_branch_output = global_out + + # The final forward applies branch_output_gain after local/global mixing. + # Scale the global branch by global_gain / output_gain here so its final + # effective gain is branch_global_output_gain while the local scaffold keeps + # branch_output_gain. The shadow mimic path still consumes raw global_out. + output_gain = self.branch_output_gain.to(dtype=hidden_states.dtype, device=hidden_states.device) + global_gain = self.branch_global_output_gain.to(dtype=hidden_states.dtype, device=hidden_states.device) + linear_out = (global_gain / torch.clamp(output_gain, min=1e-6)) * global_out + if self.local_window_enabled: + local_out = self._local_window_attention_output( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + self.last_local_window_output = local_out.detach() + local_fraction = self._local_window_fraction(dtype=hidden_states.dtype, device=hidden_states.device) + linear_out = (1.0 - local_fraction) * linear_out + local_fraction * local_out.to(dtype=hidden_states.dtype) + _debug_assert_finite("linear_branch_mix", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_branch_mix", linear_out, self.layer_idx) + return linear_out + + def 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, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + branch_past_key_values: Optional[QGRBranchCache] = None, + branch_use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if ( + self.training + and bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + and not torch.is_grad_enabled() + ): + with torch.enable_grad(): + return self.forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + fast_full_replacement = bool( + self.hybrid_enabled + and self.hybrid_replacement_mode in {"full", "replace", "linear"} + and bool(getattr(self.config, "hybrid_skip_gqa_in_full_replacement", False)) + and not (self.training and bool(getattr(self.config, "hybrid_attention_transfer_pass_gqa", False))) + ) + if fast_full_replacement: + linear_out = self._linear_attention_output( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + output_attentions=output_attentions, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + global_branch_out = getattr(self, "_last_global_branch_output", None) + linear_out = self.distill_sum(linear_out) + _debug_assert_finite("linear_distill_sum", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_distill_sum", linear_out, self.layer_idx) + gain = self.branch_output_gain.to(dtype=linear_out.dtype, device=linear_out.device) + linear_out = gain * linear_out + _debug_assert_finite("linear_output_gain", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_output_gain", linear_out, self.layer_idx) + self.last_pre_channel_output = linear_out.detach() + channel_gain = self.branch_output_channel_gain.to(dtype=linear_out.dtype, device=linear_out.device) + linear_out = linear_out * channel_gain.view(1, 1, -1) + _debug_assert_finite("linear_channel_gain", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_channel_gain", linear_out, self.layer_idx) + linear_out = self._apply_branch_output_adapter(linear_out) + _debug_assert_finite("linear_output_adapter", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_output_adapter", linear_out, self.layer_idx) + self.last_replacement_output = linear_out.detach() + self.last_linear_output = linear_out + self.last_gqa_output = None + self.last_global_linear_output = None + if ( + global_branch_out is not None + and self.local_window_enabled + and bool(getattr(self.config, "hybrid_mimic_global_branch_when_local", False)) + ): + self.last_global_linear_output = None + return linear_out.to(dtype=hidden_states.dtype), None, None + + gqa_out, attn_weights, present_key_value = super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + if not self.hybrid_enabled: + return gqa_out, attn_weights, present_key_value + + eval_mode = "" + if not self.training: + eval_mode = str(getattr(self.config, "hybrid_eval_mode", "") or "").strip().lower() + if eval_mode == "gqa_only": + return gqa_out, attn_weights, present_key_value + + mimic_return_gqa = self.training and bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + if ( + self.training + and bool(getattr(self.config, "hybrid_attention_transfer_pass_gqa", False)) + and not mimic_return_gqa + ): + self.last_gqa_output = gqa_out.detach() + self.last_replacement_output = None + self.last_linear_output = None + self.last_global_linear_output = None + self.last_quasar_output = None + self.last_raven_output = None + self.last_gla_output = None + return gqa_out, attn_weights, present_key_value + + # Safeguard to completely bypass the hybrid branch when it is gated out + # This prevents NaN propagation (0.0 * NaN = NaN) from uninitialized or unstable Triton kernels + forced_eval = eval_mode in {"quasar_forced", "raven_forced", "gla_forced", "mixed_forced"} + alpha_bypass_enabled = bool(getattr(self.config, "hybrid_alpha_zero_bypass", False)) + if alpha_bypass_enabled and float(self.replace_alpha_raw.detach().cpu()) < -13.8 and not forced_eval and not mimic_return_gqa: + return gqa_out, attn_weights, present_key_value + + linear_out = self._linear_attention_output( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + output_attentions=output_attentions, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + global_branch_out = getattr(self, "_last_global_branch_output", None) + if self.training: + self.distill_sum._distill_teacher = gqa_out.detach() + linear_out = self.distill_sum(linear_out) + _debug_assert_finite("linear_distill_sum", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_distill_sum", linear_out, self.layer_idx) + gain = self.branch_output_gain.to(dtype=linear_out.dtype, device=linear_out.device) + linear_out = gain * linear_out + _debug_assert_finite("linear_output_gain", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_output_gain", linear_out, self.layer_idx) + self.last_pre_channel_output = linear_out.detach() + channel_gain = self.branch_output_channel_gain.to(dtype=linear_out.dtype, device=linear_out.device) + linear_out = linear_out * channel_gain.view(1, 1, -1) + _debug_assert_finite("linear_channel_gain", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_channel_gain", linear_out, self.layer_idx) + linear_out = self._apply_branch_output_adapter(linear_out) + _debug_assert_finite("linear_output_adapter", linear_out, self.layer_idx) + linear_out = _sanitize_hybrid_tensor("linear_output_adapter", linear_out, self.layer_idx) + + mimic_out = linear_out + if ( + global_branch_out is not None + and self.local_window_enabled + and bool(getattr(self.config, "hybrid_mimic_global_branch_when_local", False)) + ): + global_mimic_out = self.distill_sum(global_branch_out) + _debug_assert_finite("global_mimic_distill_sum", global_mimic_out, self.layer_idx) + global_mimic_out = _sanitize_hybrid_tensor("global_mimic_distill_sum", global_mimic_out, self.layer_idx) + global_gain = self.branch_global_output_gain.to(dtype=global_mimic_out.dtype, device=global_mimic_out.device) + global_mimic_out = global_gain * global_mimic_out + _debug_assert_finite("global_mimic_output_gain", global_mimic_out, self.layer_idx) + global_mimic_out = _sanitize_hybrid_tensor("global_mimic_output_gain", global_mimic_out, self.layer_idx) + self.last_global_pre_channel_output = global_mimic_out.detach() + global_mimic_out = global_mimic_out * channel_gain.view(1, 1, -1) + _debug_assert_finite("global_mimic_channel_gain", global_mimic_out, self.layer_idx) + global_mimic_out = _sanitize_hybrid_tensor("global_mimic_channel_gain", global_mimic_out, self.layer_idx) + global_mimic_out = self._apply_branch_output_adapter(global_mimic_out) + _debug_assert_finite("global_mimic_output_adapter", global_mimic_out, self.layer_idx) + global_mimic_out = _sanitize_hybrid_tensor("global_mimic_output_adapter", global_mimic_out, self.layer_idx) + mimic_out = global_mimic_out + self.last_global_linear_output = global_mimic_out.detach() + + self.last_gqa_output = gqa_out.detach() + self.last_replacement_output = linear_out.detach() + self.last_linear_output = mimic_out + + if mimic_return_gqa: + return gqa_out, attn_weights, present_key_value + + if forced_eval: + return linear_out.to(dtype=gqa_out.dtype), attn_weights, present_key_value + + if self.hybrid_replacement_mode in {"full", "replace", "linear"}: + return linear_out.to(dtype=gqa_out.dtype), attn_weights, present_key_value + + alpha = torch.sigmoid(self.replace_alpha_raw).to(dtype=gqa_out.dtype, device=gqa_out.device) + linear_out = linear_out.to(dtype=gqa_out.dtype) + attn_output = gqa_out + alpha * linear_out + attn_output = _sanitize_hybrid_tensor("gated_hybrid_output", attn_output, self.layer_idx) + return attn_output, attn_weights, present_key_value + + +ATTENTION_CLASSES = { + "eager": QuasarLongAttention, + "flash_attention_2": QuasarLongFlashAttention2, + "sdpa": QuasarLongHybridReplacementSdpaAttention, +} + + +class QuasarLongMTPLayer(nn.Module): + def __init__(self, config: QuasarLongConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.input_layernorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.enorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.post_attention_layernorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attention = ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + self.mlp = QuasarLongSparseMoeBlock(config, layer_idx) + + self.hnorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.final_layernorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_embeds, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + def custom_mtp_attention(input_embeds_t, hidden_states_t, past_key_value_t): + input_embeds_norm = self.enorm(input_embeds_t) + hidden_states_norm = self.hnorm(hidden_states_t) + h = self.eh_proj(torch.cat([input_embeds_norm, hidden_states_norm], dim=-1)) + res = h + h_normed = self.input_layernorm(h) + + h_attn, attn_w, pres_kv = self.attention( + hidden_states=h_normed, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value_t, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + use_cache=use_cache, + ) + h_out = res + h_attn + return h_out, attn_w, pres_kv + + is_ckpt_enabled = self.training and bool(getattr(self.config, "gradient_checkpointing", False)) + if is_ckpt_enabled: + hidden_states, self_attn_weights, present_key_value = torch.utils.checkpoint.checkpoint( + custom_mtp_attention, + input_embeds, + hidden_states, + past_key_value, + use_reentrant=False, + determinism_check="none", + ) + else: + hidden_states, self_attn_weights, present_key_value = custom_mtp_attention( + input_embeds, + hidden_states, + past_key_value, + ) + + # Fully Connected (executed outside checkpoint to prevent CheckpointError in dynamic routing) + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + hidden_states = residual + hidden_states.to(residual.device) + hidden_states = self.final_layernorm(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + + +class QuasarLongDecoderLayer(nn.Module): + def __init__(self, config: QuasarLongConfig, layer_idx: int): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + self.attention = ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + + self.mlp = ( + QuasarLongSparseMoeBlock(config, layer_idx) + if (config.num_experts is not None and layer_idx >= config.first_k_dense_replace) + else QuasarLongMLP(config=config, intermediate_size=config.intermediate_size) + ) + self.input_layernorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # ── Looped-Transformer input-injection gate ────────────────────────── + # logit(-6.907) ≈ 0.001 gate at step-0, conservative while the + # looped path adapts on top of a pretrained checkpoint. + # Mirrors HybridBlock.injection_gate in quasar_rope.py. + if getattr(config, "use_looped_injection", False): + self.injection_gate = nn.Parameter(torch.tensor([-6.907])) + num_loops = max(1, int(getattr(config, "num_loops", 1))) + self.injection_gate.register_hook(lambda g: g / float(num_loops)) + else: + self.register_parameter("injection_gate", None) + + # Parcae-style loop stabilizer. This is initialized as a near-identity + # transition so pretrained checkpoints are not shocked when enabled. + if getattr(config, "use_parcae_loop_stabilizer", False): + self.parcae_decay_raw = nn.Parameter(torch.tensor([-6.907])) + self.parcae_anchor_gate = nn.Parameter(torch.tensor([-6.907])) + num_loops = max(1, int(getattr(config, "num_loops", 1))) + self.parcae_decay_raw.register_hook(lambda g: g / float(num_loops)) + self.parcae_anchor_gate.register_hook(lambda g: g / float(num_loops)) + else: + self.register_parameter("parcae_decay_raw", None) + self.register_parameter("parcae_anchor_gate", None) + + # ── Engram: static N-gram conditional memory ───────────────────────── + # Attach only to the layer indices listed in config.engram_layers. + # Falls back gracefully when engram.py is unavailable. + _engram_layers = list(getattr(config, "engram_layers", [])) + if _ENGRAM_AVAILABLE and EngramModule is not None and layer_idx in _engram_layers: + self.engram: Optional[nn.Module] = EngramModule( + vocab_size=config.vocab_size, + d_model=config.hidden_size, + d_mem=getattr(config, "engram_dim", config.hidden_size // 4), + num_heads=getattr(config, "engram_num_heads", 8), + ngram_orders=list(getattr(config, "engram_ngram_orders", [2, 3])), + target_slots=getattr(config, "engram_slots", 2_000_000), + n_layers=config.num_hidden_layers, + ) + self.engram.triton_training = bool(getattr(config, "engram_triton_training", False)) + # Mark so _init_weights skips re-initializing internal Engram params + for m in self.engram.modules(): + m._skip_quasar_hf_init = True + else: + self.engram = None + + self._engram_residual_scale = float(getattr(config, "engram_residual_scale", 0.01)) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + input_ids: Optional[torch.LongTensor] = None, # for Engram N-gram lookup + injection_P: Optional[torch.Tensor] = None, # looped-injection anchor + branch_past_key_values: Optional[QGRBranchCache] = None, + branch_use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states: (batch, seq_len, embed_dim) + input_ids: (batch, seq_len) – raw token IDs, optional; required only when + an EngramModule is attached to this layer. + injection_P: optional anchor embedding for looped-injection mixing. + (All other args identical to the standard QuasarLongDecoderLayer.) + """ + def custom_attention(h, injection_P_t, input_ids_t, past_key_value_t): + # ── Parcae-style stable recurrence: h' = decay * h + gate * P ── + if ( + injection_P_t is not None + and self.parcae_decay_raw is not None + and self.parcae_anchor_gate is not None + ): + decay = torch.exp(-F.softplus(self.parcae_decay_raw)).to(dtype=h.dtype, device=h.device) + anchor_gate = torch.sigmoid(self.parcae_anchor_gate).to(dtype=h.dtype, device=h.device) + h = decay * h + anchor_gate * injection_P_t + + # ── Looped-injection: blend residual stream with initial embeddings ── + if injection_P_t is not None and self.injection_gate is not None: + h = h + torch.sigmoid(self.injection_gate) * injection_P_t + + # ── Engram: add static N-gram memory signal before attention ───────── + if self.engram is not None and input_ids_t is not None: + engram_out, _alpha = self.engram(input_ids_t, h) + h = h + self._engram_residual_scale * engram_out + + residual_attn = h + h_normed = self.input_layernorm(h) + + # Self Attention + h_attn, attn_w, pres_kv = self.attention( + hidden_states=h_normed, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value_t, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + use_cache=use_cache, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + h_out = residual_attn + h_attn + return h_out, attn_w, pres_kv + + base_no_grad = self.training and bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + with torch.no_grad() if base_no_grad else nullcontext(): + is_ckpt_enabled = self.training and bool(getattr(self.config, "gradient_checkpointing", False)) + if is_ckpt_enabled: + hidden_states, self_attn_weights, present_key_value = torch.utils.checkpoint.checkpoint( + custom_attention, + hidden_states, + injection_P, + input_ids, + past_key_value, + use_reentrant=False, + determinism_check="none", + ) + else: + hidden_states, self_attn_weights, present_key_value = custom_attention( + hidden_states, + injection_P, + input_ids, + past_key_value, + ) + + # Fully Connected (executed outside checkpoint to prevent CheckpointError in dynamic routing) + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + hidden_states = residual + hidden_states.to(residual.device) + + outputs = (hidden_states,) + + if self.training and ( + bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + or bool(getattr(self.config, "hybrid_attention_collect_branch_loss", False)) + ): + distill_clip = float(getattr(self.config, "branch_mimic_clip", 80.0)) + all_branch_names = ("quasar", "raven", "gla", "mixed") + active_branch_set = set(getattr(self.config, "branch_mimic_branches", all_branch_names)) + branch_names = tuple(name for name in all_branch_names if name in active_branch_set) + branch_attrs = tuple( + item for item in ( + ("quasar", "last_quasar_output"), + ("raven", "last_raven_output"), + ("gla", "last_gla_output"), + ("mixed", "last_linear_output"), + ) + if item[0] in active_branch_set + ) + branch_loss = hidden_states.new_zeros((), dtype=torch.float32) + branch_loss_sums = {name: 0.0 for name in all_branch_names} + branch_cos_sums = {name: 0.0 for name in all_branch_names} + branch_rel_mse_sums = {name: 0.0 for name in all_branch_names} + branch_loss_counts = {name: 0 for name in all_branch_names} + skipped_distill = {name: 0 for name in all_branch_names} + distill_count = 0 + detailed_branch_stats = bool(getattr(self.config, "branch_mimic_detailed_stats", False)) + sanitize_checks = False + gqa_t = getattr(self.attention, "last_gqa_output", None) + if gqa_t is not None: + target = gqa_t.float().detach().clamp(-distill_clip, distill_clip) + if detailed_branch_stats: + target_flat = target.reshape(-1) + target_energy = torch.mean(target_flat * target_flat).clamp_min(1e-8) + for branch_name, attr_name in branch_attrs: + branch_s = getattr(self.attention, attr_name, None) + if branch_s is None: + continue + if sanitize_checks and not torch.isfinite(branch_s).all(): + skipped_distill[branch_name] += 1 + continue + pred = branch_s.float().clamp(-distill_clip, distill_clip) + loss_i = F.smooth_l1_loss(pred, target) + if sanitize_checks and not torch.isfinite(loss_i): + skipped_distill[branch_name] += 1 + continue + branch_loss = branch_loss + loss_i + if detailed_branch_stats: + pred_flat = pred.reshape(-1) + mse_i = torch.mean((pred_flat - target_flat) ** 2) + cos_i = F.cosine_similarity(pred_flat, target_flat, dim=0) + branch_loss_sums[branch_name] += float(loss_i.detach().item()) + branch_cos_sums[branch_name] += float(cos_i.detach().item()) if torch.isfinite(cos_i) else 0.0 + branch_rel_mse_sums[branch_name] += float((mse_i / target_energy).detach().item()) if torch.isfinite(mse_i) else 0.0 + branch_loss_counts[branch_name] += 1 + distill_count += 1 + if distill_count > 0: + branch_loss = branch_loss / distill_count + outputs += ( + branch_loss, + { + "branch_loss_sums": branch_loss_sums, + "branch_cos_sums": branch_cos_sums, + "branch_rel_mse_sums": branch_rel_mse_sums, + "branch_loss_counts": branch_loss_counts, + "skipped_distill": skipped_distill, + "distill_count": distill_count, + }, + ) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + + +QUASAR_LONG_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`QuasarLongConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare QuasarLong Model outputting raw hidden-states without any specific head on top.", + QUASAR_LONG_START_DOCSTRING, +) +class QuasarLongPreTrainedModel(PreTrainedModel): + config_class = QuasarLongConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["QuasarLongDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + # 1. Let super().from_pretrained load and instantiate the model normally + model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + + # 2. Check if we need to fuse MoE experts from separate parameters + import os + from safetensors.torch import load_file + from huggingface_hub import snapshot_download + + print(f"[FUSION LOADER] Post-loading MoE expert check/fusion for {pretrained_model_name_or_path}...", flush=True) + try: + repo_path = snapshot_download(pretrained_model_name_or_path, allow_patterns=["*.safetensors", "*.json"]) + except Exception as e: + repo_path = str(pretrained_model_name_or_path) + + files = sorted([os.path.join(repo_path, f) for f in os.listdir(repo_path) if f.endswith(".safetensors")]) + if files: + print(f"[FUSION LOADER] Analyzing safetensors for separate MoE weights in {repo_path}...", flush=True) + expert_weights = {} + has_unfused_experts = False + + for f in files: + sd = load_file(f) + for k, weight in sd.items(): + if "mlp.experts." in k: + has_unfused_experts = True + parts = k.split(".") + if "layers" in parts and "experts" in parts: + layer_idx = int(parts[parts.index("layers") + 1]) + expert_idx = int(parts[parts.index("experts") + 1]) + proj_name = parts[parts.index("experts") + 2] + expert_weights[(layer_idx, expert_idx, proj_name)] = weight + + if has_unfused_experts: + print("[FUSION LOADER] Separate experts detected! Fusing in-flight...", flush=True) + fused_sd = {} + layer_indexes = sorted(list(set(k[0] for k in expert_weights.keys()))) + for l_idx in layer_indexes: + exp_indexes = sorted(list(set(k[1] for k in expert_weights.keys() if k[0] == l_idx))) + num_exp = len(exp_indexes) + if num_exp == 0: + continue + + print(f" [FUSION LOADER] Fusing {num_exp} experts in layer {l_idx}...", flush=True) + gate_list = [] + up_list = [] + down_list = [] + for e_idx in range(num_exp): + gate_list.append(expert_weights[(l_idx, e_idx, "gate_proj")].t()) + up_list.append(expert_weights[(l_idx, e_idx, "up_proj")].t()) + down_list.append(expert_weights[(l_idx, e_idx, "down_proj")].t()) + + gate_stacked = torch.stack(gate_list) + up_stacked = torch.stack(up_list) + down_stacked = torch.stack(down_list) + + # Convert to the model's active dtype. During HF low-memory + # loading, parameters may still live on the meta device; in + # that case creating fused tensors on meta and calling a + # normal load_state_dict is a no-op, leaving MoE experts + # randomly materialized later. Keep real CPU tensors and + # assign them into the module below. + target_dtype = model.dtype + target_device = next(model.parameters()).device + if target_device.type == "meta": + target_device = torch.device("cpu") + + fused_sd[f"model.layers.{l_idx}.mlp.experts_w12"] = torch.cat([gate_stacked, up_stacked], dim=-1).to(device=target_device, dtype=target_dtype) + fused_sd[f"model.layers.{l_idx}.mlp.experts_w3"] = down_stacked.to(device=target_device, dtype=target_dtype) + + print("[FUSION LOADER] Applying fused weights to the initialized model...", flush=True) + info = model.load_state_dict(fused_sd, strict=False, assign=True) + print(f"[FUSION LOADER] Post-load fusion complete! Missing: {len(info.missing_keys)}, Unexpected: {len(info.unexpected_keys)}", flush=True) + else: + print("[FUSION LOADER] Checkpoint already contains fused weights, skipping post-load fusion.", flush=True) + else: + print("[FUSION LOADER] No safetensors files found, skipping post-load fusion.", flush=True) + + return model + + def _init_weights(self, module): + if getattr(module, "_skip_quasar_hf_init", False): + return + direct_params = list(module.parameters(recurse=False)) + direct_buffers = [buffer for buffer in module.buffers(recurse=False) if buffer is not None] + if direct_params or direct_buffers: + if all(getattr(param, "_is_hf_initialized", False) for param in direct_params) and all( + getattr(buffer, "_is_hf_initialized", False) for buffer in direct_buffers + ): + module._is_hf_initialized = True + return + if not hasattr(self, "_init_count"): + self._init_count = 0 + self._init_count += 1 + if self._init_count % 1000 == 0: + print(f" [MODEL INIT] Initializing module weights... ({self._init_count} modules processed)") + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, QuasarLongHybridReplacementSdpaAttention) and module.hybrid_enabled: + module.replace_alpha_raw.data.fill_(float(getattr(self.config, "hybrid_alpha_init", -15.0))) + module.branch_mix_logits.data.zero_() + + +QUASAR_LONG_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare QuasarLong Model outputting raw hidden-states without any specific head on top.", + QUASAR_LONG_START_DOCSTRING, +) +class QuasarLongModel(QuasarLongPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`QuasarLongDecoderLayer`] + + Args: + config: QuasarLongConfig + """ + + def __init__(self, config: QuasarLongConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.num_nextn_predict_layers = config.num_nextn_predict_layers + + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = [] + for layer_idx in range(config.num_hidden_layers + config.num_nextn_predict_layers): + if layer_idx % 1 == 0: # Print every layer for visibility + print(f"[MODEL INIT] Building layer {layer_idx}/{config.num_hidden_layers + config.num_nextn_predict_layers-1}...") + layer_cls = QuasarLongDecoderLayer if layer_idx < config.num_hidden_layers else QuasarLongMTPLayer + self.layers.append(layer_cls(config, layer_idx)) + + self.layers = nn.ModuleList(self.layers) + + self._use_sdpa = config._attn_implementation == "sdpa" + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.norm = QuasarLongRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = QuasarLongRotaryEmbedding(config=config) + self.gradient_checkpointing = False + # Initialize weights and apply final processing + print("[MODEL INIT] Finished building layers. Starting weight initialization (post_init)... this can take a few minutes for 20B models.") + self.post_init() + print("[MODEL INIT] Weight initialization complete.") + + def reset_hybrid_branch_parameters(self) -> None: + for layer in self.layers: + injection_gate = getattr(layer, "injection_gate", None) + if injection_gate is not None: + with torch.no_grad(): + injection_gate.fill_(-6.907) + attention = getattr(layer, "attention", None) + reset = getattr(attention, "reset_hybrid_branch_parameters", None) + if callable(reset): + reset() + if hasattr(layer, "engram") and layer.engram is not None: + layer.engram._init_weights() + + def get_input_embeddings(self): + return self.word_embeddings + + def set_input_embeddings(self, value): + self.word_embeddings = value + + @add_start_docstrings_to_model_forward(QUASAR_LONG_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + branch_past_key_values: Optional[QGRBranchCache] = None, + branch_use_cache: bool = False, + **kwargs, + ) -> Union[Tuple, MoeV2ModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers." + ) + use_cache = False + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if branch_use_cache and branch_past_key_values is None: + branch_past_key_values = QGRBranchCache(seen_tokens=past_seen_tokens) + if branch_use_cache and past_seen_tokens == 0 and branch_past_key_values is not None: + past_seen_tokens = int(branch_past_key_values.get_seq_length()) + + if position_ids is None: + position_ids = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + position_ids = position_ids.unsqueeze(0) + + if self._use_flash_attention_2: + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self._use_sdpa and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_seen_tokens, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, (batch_size, seq_length), inputs_embeds, past_seen_tokens + ) + + # embed positions + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + next_decoder_cache = None + layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers + mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None + + if os.environ.get("LOCAL_RANK", "0") == "0" and getattr(self, "_model_forward_debug", 0) < 1: + self._model_forward_debug = 1 + print(f"[DEBUG RANK 0] QuasarLongModel.forward started: seq_len={seq_length}", flush=True) + + # ── Looped-Transformer: anchor embedding for injection mixing ───────── + num_loops = max(1, int(getattr(self.config, "num_loops", 1))) + use_looped_injection = bool(getattr(self.config, "use_looped_injection", False)) + use_parcae_loop_stabilizer = bool(getattr(self.config, "use_parcae_loop_stabilizer", False)) + collect_branch_mimic = self.training and ( + bool(getattr(self.config, "hybrid_attention_mimic_return_gqa", False)) + or bool(getattr(self.config, "hybrid_attention_collect_branch_loss", False)) + ) + branch_mimic_loss_accum = hidden_states.new_zeros((), dtype=torch.float32) + branch_mimic_stats = None + branch_mimic_count = 0 + if collect_branch_mimic: + branch_mimic_stats = { + "branch_loss_sums": {"quasar": 0.0, "raven": 0.0, "gla": 0.0, "mixed": 0.0}, + "branch_cos_sums": {"quasar": 0.0, "raven": 0.0, "gla": 0.0, "mixed": 0.0}, + "branch_rel_mse_sums": {"quasar": 0.0, "raven": 0.0, "gla": 0.0, "mixed": 0.0}, + "branch_loss_counts": {"quasar": 0, "raven": 0, "gla": 0, "mixed": 0}, + "skipped_distill": {"quasar": 0, "raven": 0, "gla": 0, "mixed": 0}, + "distill_count": 0, + } + # P is kept as the initial embedding; each layer can blend it back in. + injection_anchor = hidden_states if (use_looped_injection or use_parcae_loop_stabilizer) else None + + for _loop_idx in range(num_loops): + for decoder_layer in layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + # Bypassed full layer checkpointing to use layer-level selective checkpointing + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + position_embeddings=position_embeddings, + input_ids=input_ids, + injection_P=injection_anchor, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + position_embeddings=position_embeddings, + input_ids=input_ids, + injection_P=injection_anchor, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + ) + hidden_states = layer_outputs[0] + + if collect_branch_mimic: + layer_branch_loss = layer_outputs[1] + layer_stats = layer_outputs[2] + layer_count = int(layer_stats.get("distill_count", 0)) + if layer_count > 0: + branch_mimic_loss_accum = branch_mimic_loss_accum + layer_branch_loss + branch_mimic_count += 1 + branch_mimic_stats["distill_count"] += layer_count + for stat_name in ( + "branch_loss_sums", + "branch_cos_sums", + "branch_rel_mse_sums", + "branch_loss_counts", + "skipped_distill", + ): + for branch_name, value in layer_stats.get(stat_name, {}).items(): + branch_mimic_stats[stat_name][branch_name] += value + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + + + hidden_states = self.norm(hidden_states) + main_hidden_states = hidden_states + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (main_hidden_states,) + + mtp_hidden_states = None + + if mtp_layers: + for decoder_layer in mtp_layers: + input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1) + inputs_embeds = self.word_embeddings(input_ids) + + if self.gradient_checkpointing and self.training: + # Bypassed full layer checkpointing to use layer-level selective checkpointing + layer_outputs = decoder_layer( + inputs_embeds, + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + else: + layer_outputs = decoder_layer( + inputs_embeds, + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + if mtp_hidden_states is None: + mtp_hidden_states = [] + hidden_states = layer_outputs[0] + mtp_hidden_states.append(hidden_states) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + branch_mimic_loss = None + if collect_branch_mimic: + branch_mimic_loss = ( + branch_mimic_loss_accum / branch_mimic_count + if branch_mimic_count > 0 + else branch_mimic_loss_accum + ) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache + if not return_dict: + return tuple( + v + for v in [ + main_hidden_states, + next_cache, + branch_past_key_values if branch_use_cache else None, + all_hidden_states, + all_self_attns, + all_router_logits, + ] + if v is not None + ) + return MoeV2ModelOutputWithPast( + last_hidden_state=main_hidden_states, + past_key_values=next_cache, + branch_past_key_values=branch_past_key_values if branch_use_cache else None, + hidden_states=all_hidden_states, + mtp_hidden_states=mtp_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + branch_mimic_loss=branch_mimic_loss, + branch_mimic_stats=branch_mimic_stats, + ) + + +class QuasarLongForCausalLM(QuasarLongPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: QuasarLongConfig): + super().__init__(config) + self.model = QuasarLongModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.num_nextn_predict_layers = config.num_nextn_predict_layers + self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor + + # Initialize weights and apply final processing + self.post_init() + + def reset_hybrid_branch_parameters(self) -> None: + self.model.reset_hybrid_branch_parameters() + + def get_input_embeddings(self): + return self.model.word_embeddings + + def set_input_embeddings(self, value): + self.model.word_embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(QUASAR_LONG_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=MoEV2CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + logit_indices: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + branch_past_key_values: Optional[QGRBranchCache] = None, + branch_use_cache: bool = False, + **kwargs, + ) -> Union[Tuple, MoEV2CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer + + >>> model = QuasarLongForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + fast_ce_labels = kwargs.pop("fast_ce_labels", None) + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + branch_past_key_values=branch_past_key_values, + branch_use_cache=branch_use_cache, + **kwargs, + ) + + loss = None + all_mtp_loss = None + aux_loss = None + hidden_states = outputs[0] + skip_logits = ( + self.training + and labels is None + and bool(getattr(self.config, "branch_mimic_skip_logits", False)) + and not bool(getattr(self.config, "branch_mimic_compute_logits", True)) + ) + if fast_ce_labels is not None: + if LigerFusedLinearCrossEntropyLoss is None: + raise RuntimeError("fast_ce_labels requested but liger_kernel is not available") + if not hasattr(self, "_quasar_liger_ce"): + self._quasar_liger_ce = LigerFusedLinearCrossEntropyLoss(ignore_index=-100) + ce_hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) + ce_target = fast_ce_labels.to(device=ce_hidden.device, dtype=torch.long).reshape(-1) + loss = self._quasar_liger_ce(self.lm_head.weight, ce_hidden, ce_target) + logits = hidden_states.new_empty((hidden_states.shape[0], hidden_states.shape[1], 0), dtype=torch.float32) + elif skip_logits: + logits = hidden_states.new_empty((hidden_states.shape[0], hidden_states.shape[1], 0), dtype=torch.float32) + elif logit_indices is not None: + if labels is not None: + raise ValueError("labels are not supported with logit_indices") + if logit_indices.shape[1] > hidden_states.shape[1]: + raise ValueError( + f"logit_indices sequence length {logit_indices.shape[1]} exceeds hidden length {hidden_states.shape[1]}" + ) + selected_hidden = hidden_states[:, : logit_indices.shape[1], :] + flat_indices = logit_indices.to(device=self.lm_head.weight.device, dtype=torch.long).reshape(-1) + selected_weight = self.lm_head.weight.index_select(0, flat_indices) + selected_weight = selected_weight.view(*logit_indices.shape, selected_hidden.shape[-1]) + logits = torch.einsum("bsh,bskh->bsk", selected_hidden, selected_weight) + else: + if isinstance(logits_to_keep, int): + hidden_for_logits = hidden_states[:, -logits_to_keep:, :] if logits_to_keep > 0 else hidden_states + else: + hidden_for_logits = hidden_states[:, logits_to_keep, :] + logits = self.lm_head(hidden_for_logits) + if labels is not None: + logits = logits.float() + if logits.numel() > 0 and labels is not None and not torch.isfinite(logits).all(): + rank = os.environ.get("LOCAL_RANK", "0") + if rank == "0": + finite_mask = torch.isfinite(logits) + nonfinite_count = (~finite_mask).sum().item() + if finite_mask.any(): + # Safely extract min/max of finite elements without indexing + # Replace non-finite elements with large positive/negative values for min/max calculation + logits_for_min = torch.where(finite_mask, logits, torch.tensor(float('inf'), device=logits.device, dtype=logits.dtype)) + logits_for_max = torch.where(finite_mask, logits, torch.tensor(float('-inf'), device=logits.device, dtype=logits.dtype)) + print( + "[DEBUG RANK 0] Non-finite logits before loss: " + f"finite_min={logits_for_min.min().item():.4e} " + f"finite_max={logits_for_max.max().item():.4e} " + f"nonfinite={nonfinite_count}", + flush=True, + ) + else: + print("[DEBUG RANK 0] Non-finite logits before loss: all logits non-finite", flush=True) + + if labels is not None: + # --- LOSS DEBUG --- + if os.environ.get("LOCAL_RANK", "0") == "0" and getattr(self, "_loss_debug_count", 0) < 5: + self._loss_debug_count = getattr(self, "_loss_debug_count", 0) + 1 + print(f"[DEBUG RANK 0] Step {self._loss_debug_count}: labels[:5]={labels.reshape(-1)[:5].tolist()}, vocab={self.config.vocab_size}", flush=True) + # Check if labels are all -100 + if (labels == -100).all(): + print("[DEBUG RANK 0] WARNING: All labels are -100! Loss will be 0.", flush=True) + + loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs) + + if os.environ.get("LOCAL_RANK", "0") == "0" and getattr(self, "_loss_debug_count", 0) <= 5: + print(f"[DEBUG RANK 0] Calculated loss: {loss.item() if loss is not None else 'None'}", flush=True) + + all_mtp_logits = None + if self.num_nextn_predict_layers > 0: + mtp_hidden_states = outputs.mtp_hidden_states + shift_labels_mtp = None + keep_mtp_logits = (not self.training) or (labels is None and fast_ce_labels is None) + for i in range(self.num_nextn_predict_layers): + mtp_hidden_states = mtp_hidden_states[i] + mtp_logits = self.lm_head(mtp_hidden_states) + if keep_mtp_logits: + if all_mtp_logits is None: + all_mtp_logits = [] + all_mtp_logits.append(mtp_logits) + if labels is not None: + if shift_labels_mtp is None: + shift_labels_mtp = labels.clone() + shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100) + mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size) + mtp_loss = self.loss_function(mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs) + if loss is not None: + loss += self.mtp_loss_scaling_factor * mtp_loss + else: + loss = self.mtp_loss_scaling_factor * mtp_loss + + if all_mtp_loss is None: + all_mtp_loss = [] + all_mtp_loss.append(mtp_loss) + del mtp_logits + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoEV2CausalLMOutputWithPast( + loss=loss, + mtp_loss=all_mtp_loss, + aux_loss=aux_loss, + branch_mimic_loss=getattr(outputs, "branch_mimic_loss", None), + branch_mimic_stats=getattr(outputs, "branch_mimic_stats", None), + logits=logits, + mtp_logits=all_mtp_logits, + past_key_values=outputs.past_key_values, + branch_past_key_values=getattr(outputs, "branch_past_key_values", None), + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) diff --git a/quasar_banner.png b/quasar_banner.png new file mode 100644 index 0000000000000000000000000000000000000000..7beeb3eed80c31870616d0d633c8e9a179a04392 --- /dev/null +++ b/quasar_banner.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cdf1e83acc41d0b9a75d4ecaa20026cecb2272b01f957ffb4fc7acc03ef33932 +size 308841 diff --git a/tokenizer.json b/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..92c58654a6814e4f28badeef13dbcec05230e0bd --- /dev/null +++ b/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fdcadf59ad1db38dde175f2a82d3ec2dde15986ac1f81aef69c5cdd03afc6e1b +size 12205847 diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..87b40c3e1e454f81f55eb19abcca3de1dc10faf9 --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,18 @@ +{ + "backend": "tokenizers", + "bos_token": "<|startoftext|>", + "clean_up_tokenization_spaces": false, + "cls_token": "[CLS]", + "eos_token": "<|endoftext|>", + "fast_tokenizer": true, + "gmask_token": "[gMASK]", + "is_local": false, + "local_files_only": false, + "merges_file": null, + "model_max_length": 1000000000000000019884624838656, + "model_specific_special_tokens": { + "gmask_token": "[gMASK]" + }, + "pad_token": "<|endoftext|>", + "tokenizer_class": "TokenizersBackend" +}