add remote code + model files

.ipynb_checkpoints/modeling_sliding_window-checkpoint.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_sliding_window import SlidingWindowConfig
+from forgetting_transformer.ops.sliding_window_attention_std import sliding_window_attention_std
+
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, config: SlidingWindowConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_kv_heads = config.num_kv_heads or self.num_heads
+        self.window_size = config.window_size
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+        
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
+        B, T, H = hidden_states.shape
+        
+        # Project
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        
+        # Reshape
+        q = q.view(B, T, self.num_heads, self.head_dim)
+        k = k.view(B, T, self.num_kv_heads, self.head_dim)
+        v = v.view(B, T, self.num_kv_heads, self.head_dim)
+        
+        # Sliding window attention
+        attn_output = sliding_window_attention_std(
+            q, k, v,
+            head_first=False,
+            window_size=self.window_size,
+        )
+        
+        # Output projection
+        attn_output = attn_output.reshape(B, T, self.hidden_size)
+        output = self.o_proj(attn_output)
+        
+        return output, None
+
+
+class SlidingWindowMLP(nn.Module):
+    def __init__(self, config: SlidingWindowConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        
+        self.gate_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.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.act_fn = nn.SiLU()
+        
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SlidingWindowDecoderLayer(nn.Module):
+    def __init__(self, config: SlidingWindowConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        
+        self.attn = SlidingWindowAttention(config, layer_idx)
+        self.mlp = SlidingWindowMLP(config)
+        
+        self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
+        # Attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.attn(hidden_states, attention_mask)
+        hidden_states = residual + hidden_states
+        
+        # MLP
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        return hidden_states, None
+
+
+class SlidingWindowModel(PreTrainedModel):
+    config_class = SlidingWindowConfig
+    _no_split_modules = ["SlidingWindowDecoderLayer"]  # ← 关键修复1
+    
+    def __init__(self, config: SlidingWindowConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([
+            SlidingWindowDecoderLayer(config, layer_idx)
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        
+        self.gradient_checkpointing = False
+        self.post_init()
+        
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        hidden_states = self.embed_tokens(input_ids)
+        
+        for decoder_layer in self.layers:
+            hidden_states, _ = decoder_layer(hidden_states, attention_mask)
+        
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+
+class SlidingWindowForCausalLM(PreTrainedModel):
+    config_class = SlidingWindowConfig
+    _tied_weights_keys = ["lm_head.weight"]
+    _no_split_modules = ["SlidingWindowDecoderLayer"]  # ← 关键修复2
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SlidingWindowModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+    
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = 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 forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        **kwargs
+    ):
+        hidden_states = self.model(input_ids, attention_mask)
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            # Return per-token loss with shape [B, T-1]
+            loss_fct = nn.CrossEntropyLoss(reduction='none')
+            loss = loss_fct(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1)
+            )
+            # Reshape to [B, T-1]
+            B, T = shift_logits.size(0), shift_logits.size(1)
+            loss = loss.view(B, T)
+            
+            # Pad last position to make shape [B, T] instead of [B, T-1]
+            loss = F.pad(loss, (0, 1), value=0.0)
+        
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+        )
+    
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}

__init__.py

@@ -0,0 +1 @@
+# for HF remote code

configuration_sliding_window.py

@@ -0,0 +1,70 @@
+from transformers import PretrainedConfig
+
+
+class SlidingWindowConfig(PretrainedConfig):
+    model_type = "sliding_window"
+    
+    def __init__(
+        self,
+        vocab_size=50304,
+        hidden_size=768,
+        intermediate_size=None,
+        hidden_ratio=4,
+        num_hidden_layers=12,
+        num_heads=12,
+        num_kv_heads=None,
+        hidden_act="swish",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        attention_bias=False,
+        fuse_norm=True,
+        fuse_cross_entropy=True,
+        use_rope=False,
+        # Sliding window specific
+        window_size=2,  # 默认2-gram
+        qk_norm=False,
+        qk_norm_share_param_across_head=False,
+        use_k_shift=False,
+        use_v_shift=False,
+        elementwise_affine=True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size or hidden_ratio * hidden_size
+        self.hidden_ratio = hidden_ratio
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.norm_eps = norm_eps
+        self.use_cache = use_cache
+        self.tie_word_embeddings = tie_word_embeddings
+        self.attention_bias = attention_bias
+        self.fuse_norm = fuse_norm
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.use_rope = use_rope
+        
+        # Sliding window
+        self.window_size = window_size
+        
+        self.qk_norm = qk_norm
+        self.qk_norm_share_param_across_head = qk_norm_share_param_across_head
+        self.use_k_shift = use_k_shift
+        self.use_v_shift = use_v_shift
+        self.elementwise_affine = elementwise_affine
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )

modeling_sliding_window.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_sliding_window import SlidingWindowConfig
+from forgetting_transformer.ops.sliding_window_attention_std import sliding_window_attention_std
+
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, config: SlidingWindowConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_kv_heads = config.num_kv_heads or self.num_heads
+        self.window_size = config.window_size
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+        
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
+        B, T, H = hidden_states.shape
+        
+        # Project
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        
+        # Reshape
+        q = q.view(B, T, self.num_heads, self.head_dim)
+        k = k.view(B, T, self.num_kv_heads, self.head_dim)
+        v = v.view(B, T, self.num_kv_heads, self.head_dim)
+        
+        # Sliding window attention
+        attn_output = sliding_window_attention_std(
+            q, k, v,
+            head_first=False,
+            window_size=self.window_size,
+        )
+        
+        # Output projection
+        attn_output = attn_output.reshape(B, T, self.hidden_size)
+        output = self.o_proj(attn_output)
+        
+        return output, None
+
+
+class SlidingWindowMLP(nn.Module):
+    def __init__(self, config: SlidingWindowConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        
+        self.gate_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.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.act_fn = nn.SiLU()
+        
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SlidingWindowDecoderLayer(nn.Module):
+    def __init__(self, config: SlidingWindowConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        
+        self.attn = SlidingWindowAttention(config, layer_idx)
+        self.mlp = SlidingWindowMLP(config)
+        
+        self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
+        # Attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.attn(hidden_states, attention_mask)
+        hidden_states = residual + hidden_states
+        
+        # MLP
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        return hidden_states, None
+
+
+class SlidingWindowModel(PreTrainedModel):
+    config_class = SlidingWindowConfig
+    _no_split_modules = ["SlidingWindowDecoderLayer"]  # ← 关键修复1
+    
+    def __init__(self, config: SlidingWindowConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([
+            SlidingWindowDecoderLayer(config, layer_idx)
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, elementwise_affine=config.elementwise_affine)
+        
+        self.gradient_checkpointing = False
+        self.post_init()
+        
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        hidden_states = self.embed_tokens(input_ids)
+        
+        for decoder_layer in self.layers:
+            hidden_states, _ = decoder_layer(hidden_states, attention_mask)
+        
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+
+class SlidingWindowForCausalLM(PreTrainedModel):
+    config_class = SlidingWindowConfig
+    _tied_weights_keys = ["lm_head.weight"]
+    _no_split_modules = ["SlidingWindowDecoderLayer"]  # ← 关键修复2
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SlidingWindowModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+    
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = 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 forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        **kwargs
+    ):
+        hidden_states = self.model(input_ids, attention_mask)
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            # Return per-token loss with shape [B, T-1]
+            loss_fct = nn.CrossEntropyLoss(reduction='none')
+            loss = loss_fct(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1)
+            )
+            # Reshape to [B, T-1]
+            B, T = shift_logits.size(0), shift_logits.size(1)
+            loss = loss.view(B, T)
+            
+            # Pad last position to make shape [B, T] instead of [B, T-1]
+            loss = F.pad(loss, (0, 1), value=0.0)
+        
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+        )
+    
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}

ops/.ipynb_checkpoints/forgetting_attention-checkpoint.py

@@ -0,0 +1,1138 @@
+"""
+Implementation of Forgetting Attention.
+
+Our code is adapted from https://github.com/FlagOpen/FlagAttention/blob/ee91638dec6da8c00c4113d179f469e0ffcd5852/src/flag_attn/flash.py. The code is modified to implement Forgetting Attention.
+
+The original license info from FlagAttention:
+
+Copyright 2023 BAAI
+
+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 pytest
+import math
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+from typing import Optional
+
+
+__all__ = ["forgetting_attention"]
+
+
+# File flash.py
+def maybe_contiguous(x):
+    # only when the inner most dimension is contiguous can LDGSTS be used
+    # so inner-dimension contiguity is enforced.
+    return x.contiguous() if x.stride(-1) != 1 else x
+
+def rounded_multiple(a, b):
+    return (a + b - 1) // b * b
+
+# --------------------------- public API ---------------------------
+class ForgettingAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, log_fgate, seq_start, causal, sm_scale, return_log_normalizer):
+        assert causal, "Only causal attention is supported"
+        Dq, Dk, Dv = q.shape[-1], k.shape[-1], v.shape[-1]
+        assert Dq == Dk == Dv, "feature size of q, k, v should be equal"
+        assert Dk in {16, 32, 64, 128}, "We only support head dims in {16, 32, 64, 128}"
+
+        B, H, M, D = q.shape
+        if seq_start is not None:
+            has_seq_start = True
+            assert seq_start.shape == (B,)
+        else:
+            has_seq_start = False
+            seq_start = torch.zeros((B,), device=q.device, dtype=torch.long)
+        N = k.shape[2]
+        assert log_fgate.shape == (B, H, N)
+        log_fgate = log_fgate.float()
+        if has_seq_start:
+            log_fgate = log_fgate.clone()
+            # We absolutely don't want masked value to affect result. If we
+            # don't do this then it could via affecting numerical precision of
+            # cumsum
+            mask_index = (torch.arange(N, device=q.device)[None, None, :] < seq_start[:, None, None])
+            mask_index = torch.broadcast_to(mask_index, log_fgate.size())
+            log_fgate[mask_index] = 0.0
+
+        log_lambda = torch.cumsum(log_fgate, dim=-1, dtype=log_fgate.dtype).float()
+
+        Hk, Hv = k.shape[1], v.shape[1]
+        assert Hk == Hv, "num of heads in k and v should be equal"
+        assert H == Hk, "groupped query attention has not been tested. You can uncomment this if you know what you are doing."
+        assert H % Hk == 0, "number of heads in q must be a multiple of that in k & v"
+        num_groups = H // Hk
+
+        P_SEQ = N - M
+        larger_m = M > N
+        assert (not larger_m), "The key/value tensors must be longer than the query tensor"
+
+        if sm_scale is None:
+            sm_scale = 1. / math.sqrt(D)
+
+        # contiguity
+        q, k, v = maybe_contiguous(q), maybe_contiguous(k), maybe_contiguous(v)
+
+        # to work around https://github.com/openai/triton/issues/2441
+        device = torch.cuda.device_of(q)
+
+        with torch.cuda.device(device):
+
+            config = get_fwd_config(B, H, M, N, D, causal)
+            BLOCK_M, BLOCK_N, num_stages, num_warps = config
+
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+            # consider using 3d grid to avoid div & rem
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            o = torch.empty_like(q)
+            L = torch.empty((B, H, M), device=q.device, dtype=torch.float32)
+            _fwd_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale,
+                L, o,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+                B, H, M, N, P_SEQ, num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=D,
+                IS_CAUSAL=causal, LARGER_M=larger_m, HAS_SEQ_START=has_seq_start,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
+                num_warps=num_warps, num_stages=num_stages,
+            )
+
+        # autograd context maintenance
+        ctx.save_for_backward(q, k, v, o, L, log_lambda, seq_start)
+        ctx.sm_scale = sm_scale
+        ctx.causal = causal
+        ctx.has_seq_start = has_seq_start
+
+        has_extra_return = return_log_normalizer
+        if has_extra_return:
+            outs = (
+                o,
+                L if return_log_normalizer else None,
+            )
+            return outs
+        return o
+
+    @staticmethod
+    def backward(ctx, do, *ignored):
+        q, k, v, o, L, log_lambda, seq_start = ctx.saved_tensors
+        sm_scale = ctx.sm_scale
+        causal = ctx.causal
+        has_seq_start = ctx.has_seq_start
+
+        B, H, M, D = q.shape
+        N = k.shape[2]
+        Hk = k.shape[1]
+        num_groups = H // Hk
+        P_SEQ = N - M
+        larger_m = M > N
+
+        if sm_scale is None:
+            sm_scale = 1. / math.sqrt(D)
+
+        # to work around https://github.com/openai/triton/issues/2441
+        device = torch.cuda.device_of(q)
+        with torch.cuda.device(device):
+            config = get_bwd_config(B, H, M, N, D, causal)
+            BLOCK_M, BLOCK_N, num_stages, num_warps = config
+
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+
+            delta = torch.empty_like(L)
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            _bwd_preprocess[grid](
+                o, do,
+                delta,
+                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                delta.stride(0), delta.stride(1), delta.stride(2),
+                M,
+                BLOCK_M=BLOCK_M, D_HEAD=D,
+                DIVISIBLE_M=divisible_m,
+            )
+
+            # NOTE that dk & dv always have the same number of heads as q, instead of q.
+            BLOCK_M, BLOCK_N, num_stages, num_warps = get_bwd_kv_config(B, H, M, N, D, causal)
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+
+            dk = torch.empty((B, H, N, D), dtype=k.dtype, device=q.device)
+            dv = torch.empty((B, H, N, D), dtype=v.dtype, device=q.device)
+            dlog_lambda = torch.empty((B, H, N), dtype=log_lambda.dtype, device=q.device)
+            grid = (triton.cdiv(N, BLOCK_N), H, B)
+            _bwd_kv_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale, do,
+                dk, dv, dlog_lambda,
+                L, delta,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
+                dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3),
+                dlog_lambda.stride(0), dlog_lambda.stride(1), dlog_lambda.stride(2),
+                B, H, M, N, P_SEQ,
+                num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N, CAUSAL=causal,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n, HAS_SEQ_START=has_seq_start,
+                num_stages=num_stages, num_warps=num_warps,
+            )
+
+            BLOCK_M, BLOCK_N, num_stages, num_warps = get_bwd_q_config(B, H, M, N, D, causal)
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+            dq = torch.zeros_like(q)
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            _bwd_q_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale, do,
+                dq, dlog_lambda,
+                L, delta,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
+                dlog_lambda.stride(0), dlog_lambda.stride(1), dlog_lambda.stride(2),
+                B, H, M, N, P_SEQ,
+                num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
+                CAUSAL=causal, LARGER_M=larger_m, HAS_SEQ_START=has_seq_start,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
+                num_stages=num_stages, num_warps = num_warps,
+            )
+            dk = dk.reshape((B, Hk, num_groups, N, D)).sum(2)
+            dv = dv.reshape((B, Hk, num_groups, N, D)).sum(2)
+        dcumsum = torch.cumsum(dlog_lambda, dim=-1, dtype=log_lambda.dtype)
+        dlog_fgate = dlog_lambda + dcumsum[..., -1:] - dcumsum
+        dlog_fgate = dlog_fgate.float()
+        return dq, dk, dv, dlog_fgate, None, None, None, None, None, None, None
+
+
+def forgetting_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_fgate: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+):
+    """
+    A FlashAttention-based implementation of Forgetting Attention. 
+
+    Note:
+    - We recommand bfloat16/float16 for q, k, v and float32 for log_fgate. float32 for 
+      q, k, v is also supported, but the kernel will not use tensor cores if q, k, v are
+      in float32 (which would be slow).
+    - We only support seqlen_q <= seqlen_k
+    - We only support causal attention
+    - Head dimension must be in one of {16, 32, 64, 128}
+
+    Arguments:
+        - q: (batch_size, seqlen_q, num_heads, head_dim) unless head_first=True.
+        - k: (batch_size, seqlen_k, num_heads, head_dim) unless head_first=True.
+        - v: (batch_size, seqlen_k, num_heads, head_dim) unless head_first=True.
+        - log_fgate: (batch_size, seqlen_k, num_heads) unless head_first=True. 
+              This should be the **log** of the forget gates. This is typically the 
+              output of torch.nn.functional.logsigmoid.
+        - head_first: if True, the order the num_heads and seqlen_* axis of the all 
+              FloatTensor inputs and outputs should be (num_heads, seq_len_*) instead of
+              (seq_len_*, num_heads)
+        - seq_start: If not None, should be LongTensor with shape (batch_size,) 
+              and range in [0, seq_len_k). For each batch index batch_id, no attention 
+              will be allocated to tokens before the token index seq_start[batch_id]. 
+              This is useful for left-padded inputs.
+        - sm_scale: The scaling of attention scores before applying softmax. If
+              None, it defaults to (1.0 / math.sqrt(head_dim))
+
+    Returns:
+        out (torch.Tensor): (batch_size, seqlen_q, num_heads, head_dim) unless head_first=True.
+    """
+    if not head_first:
+        q, k, v = [rearrange(item, "b t h d -> b h t d") for item in (q, k, v)]
+        log_fgate = rearrange(log_fgate, "b t h -> b h t")
+    out = ForgettingAttention.apply(q, k, v, log_fgate, seq_start, True, sm_scale, False)
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    return out
+
+
+# --------------------------- Forward ---------------------------
+# NOTE: this function can be overwritten at runtime to use your custom config
+def get_fwd_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0):
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 32, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 4, 4
+    elif torch.cuda.get_device_capability() == (9, 0):
+        # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 2, 8
+    elif torch.cuda.get_device_capability() == (8, 6):
+        if not causal:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+        else: # causal
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 3, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9):
+        # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+
+@triton.jit
+def _fwd_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale,
+    L, O,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr, LARGER_M: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_m = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    log2e: tl.constexpr = 1.4426950408889634
+    loge2: tl.constexpr = 0.6931471805599453
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    O += off_z * stride_oz + off_h * stride_oh
+    L += (off_z * H + off_h) * M # l's shape is (B, H, M)
+
+    offs_m_base = tl.arange(0, BLOCK_M)
+    offs_m = start_m * BLOCK_M + offs_m_base
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m) * stride_log_lambda_n
+    o_ptrs = O + (offs_m[:, None] * stride_om + offs_k[None, :] * stride_ok) # (BLOCK_M, BLOCK_DMODEL)
+    l_ptrs = L + offs_m
+
+    # initialize pointer to m and l, fp32 for accumulators
+    m_i = tl.full([BLOCK_M], value=-float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    # load q
+    if DIVISIBLE_M:
+        q = tl.load(q_ptrs, cache_modifier=".cg")
+        log_lambda_out = tl.load(log_lambda_out_ptrs, cache_modifier=".cg")
+    else:
+        mask_m = offs_m < M
+        q = tl.load(q_ptrs, mask=mask_m[:, None], cache_modifier=".cg")
+        log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m, cache_modifier=".cg")
+
+    #Dot I trick: to place q in registers, it saves shared memory
+    # if BLOCK_DMODEL < 128:
+    #     I = tl.where(offs_k[:, None] == offs_k,
+    #                  tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 1.0, dtype=input_dtype),
+    #                  tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 0.0, dtype=input_dtype))
+    #     q = tl.dot(q, I, input_precision="ieee").to(input_dtype)
+    # else:
+    #     I = tl.where(offs_m_base[:, None] == offs_m_base,
+    #                  tl.full((BLOCK_M, BLOCK_M), 1.0, dtype=input_dtype),
+    #                  tl.full((BLOCK_M, BLOCK_M), 0.0, dtype=input_dtype))
+    #     q = tl.dot(I, q, input_precision="ieee").to(input_dtype)
+
+    # NOTE: Loop-Bound-For-N
+    # The indices in m-dimension that this block may access is in `[start_m * BLOCK_M, (start_m + 1) * BLOCK_M)`.
+    # According to the rule of causal masking, then max index in n-dimension that this block may access
+    # is `P_SEQ + (start_m + 1) * BLOCK_M`.
+    # However, the upper bound of index in n-dimension should never exceed the sequence length of k/v(`P_SEQ + N_CTX`).
+    # `P_SEQ + (start_m + 1) * BLOCK_M` may be larger than `N`.
+    # At this case, there would be illegal memory access when loading k & v tiles
+    # if mask_n is not applied for loading(only when `DIVISIBLE_N`` is true).
+    # See also https://github.com/FlagOpen/FlagAttention/pull/8
+    if IS_CAUSAL:
+        hi = tl.minimum(N, P_SEQ + (start_m + 1) * BLOCK_M)
+        if LARGER_M:
+            hi = tl.maximum(0, hi)
+    else:
+        hi = N
+
+    offs_n_init = offs_n_base
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        lo = tl.minimum(seq_start, hi)
+        lo = (lo // BLOCK_N) * BLOCK_N
+        offs_n_init += lo
+    else:
+        lo = 0
+        seq_start = 0
+
+    # loop over k, v and update accumulators
+    k_ptrs = K + (offs_k[:, None] * stride_kk + offs_n_init[None, :] * stride_kn) # (BLOCK_DMODEL, BLOCK_N)
+    v_ptrs = V + (offs_n_init[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n_init * stride_log_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+    for start_n in range(lo, hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        offs_n = start_n + offs_n_base
+
+        # -- load k, v --
+        if DIVISIBLE_N:
+            k = tl.load(k_ptrs, cache_modifier=".cg")
+            v = tl.load(v_ptrs, cache_modifier=".cg")
+            log_lambda_in = tl.load(log_lambda_in_ptrs, cache_modifier=".cg")
+        else:
+            mask_n = offs_n < N
+            k = tl.load(k_ptrs, mask=mask_n[None, :], cache_modifier=".cg")
+            v = tl.load(v_ptrs, mask=mask_n[:, None], cache_modifier=".cg")
+            log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n, cache_modifier=".cg")
+
+        # -- compute qk ---
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        s = tl.dot(q, k, input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[:, None] - log_lambda_in[None, :]
+        s += decay_bias * log2e
+
+        if not DIVISIBLE_N:
+            s = tl.where(mask_n[None, :], s, float("-inf"))
+        if IS_CAUSAL:
+            causal_mask = (P_SEQ + offs_m[:, None]) >= offs_n[None, :]
+            s = tl.where(causal_mask, s, float("-inf"))
+        if HAS_SEQ_START:
+            s = tl.where(offs_n[None, :] >= seq_start, s, float("-inf"))
+
+
+        # -- compute scaling constant ---
+        m_i_new = tl.maximum(m_i, tl.max(s, 1))
+        alpha = tl.math.exp2((m_i - m_i_new))
+        p = tl.math.exp2(s - m_i_new[:, None])
+
+        # -- compute partial sumexpn before applying dropout
+        p_sum = tl.sum(p, 1)
+
+
+        # -- scale and update acc: acc *= alpha[:, None]--
+        acc *= alpha[:, None]
+        acc += tl.dot(p.to(input_dtype), v, input_precision="ieee")
+
+        # -- update m_i and l_i --
+        l_i = l_i * alpha + p_sum
+        m_i = m_i_new
+        # update pointers
+        k_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vn
+        log_lambda_in_ptrs += BLOCK_N * stride_log_lambda_n
+
+    # write back l & o
+    if IS_CAUSAL and (LARGER_M or HAS_SEQ_START):
+        is_empty_line = (offs_m + P_SEQ) < seq_start
+        acc = tl.where(is_empty_line[:, None], 0.0, acc * (1.0 / l_i[:, None]))
+        l = tl.where(is_empty_line, float("-inf"), m_i * loge2 + tl.log(l_i))
+    else:
+        acc = acc * (1.0 / l_i[:, None])
+        l = m_i * loge2 + tl.log(l_i) # log(normalizer)
+
+
+    if DIVISIBLE_M:
+        tl.store(l_ptrs, l, cache_modifier=".cg")
+        tl.store(o_ptrs, acc.to(input_dtype), cache_modifier=".cg")
+    else:
+        tl.store(l_ptrs, l, mask=mask_m, cache_modifier=".cg")
+        tl.store(o_ptrs, acc.to(input_dtype), mask=mask_m[:, None], cache_modifier=".cg")
+
+
+# --------------------------- Backward ---------------------------
+# NOTE: this function can be overwritten at runtime to use your custom config
+def get_bwd_config(B, H, M, N, D, causal):
+    if torch.cuda.get_device_capability() == (9, 0):
+        if not causal:
+            BLOCK_M = 128 if D <= 64 else 64
+            BLOCK_N = 64
+            num_stages = 2
+            num_warps = 4
+        else:
+            BLOCK_M = 64
+            BLOCK_N = 64
+            num_stages = 3 if D <= 64 else 2
+            num_warps = 4
+    elif torch.cuda.get_device_capability() == (8, 0):
+        if not causal:
+            BLOCK_M = 128 if D <= 64 else 64
+            BLOCK_N = 64
+            num_stages = 2
+            num_warps = 4
+        else:
+            BLOCK_M = 64
+            BLOCK_N = 64
+            num_stages = 3 if D <= 64 else 2
+            num_warps = 4
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if not causal:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 8
+        else:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 1, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+def get_bwd_kv_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0): # A100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 4, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 128, 4, 8
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9): # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 128, 4, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 128, 2, 8
+    elif torch.cuda.get_device_capability() == (9, 0): # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+def get_bwd_q_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0): # A100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 4, 8
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9): # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 4, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 3, 4
+    elif torch.cuda.get_device_capability() == (9, 0): # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 4, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 2, 8
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+
+@triton.jit
+def _bwd_preprocess(
+    Out, DO,
+    Delta,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dz, stride_dh, stride_dm,
+    M,
+    BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr,
+):
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+    Out += off_z * stride_oz + off_h * stride_oh
+    DO += off_z * stride_doz + off_h * stride_doh
+    Delta += off_z * stride_dz + off_h * stride_dh
+
+    # compute (Out * Dout).sum() for vector interpretation
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_n = tl.arange(0, D_HEAD)
+
+    # load
+    o_ptrs = Out + off_m[:, None] * stride_om + off_n[None, :] * stride_ok
+    do_ptrs = DO + off_m[:, None] * stride_dom + off_n[None, :] * stride_dok
+
+    if DIVISIBLE_M:
+        o = tl.load(o_ptrs).to(tl.float32)
+        do = tl.load(do_ptrs).to(tl.float32)
+    else:
+        mask_m = off_m < M
+        o = tl.load(o_ptrs, mask=mask_m[:, None]).to(tl.float32)
+        do = tl.load(do_ptrs, mask=mask_m[:, None]).to(tl.float32)
+
+    # compute
+    delta = tl.sum(o * do, axis=1)
+
+    # write-back
+    d_ptrs = Delta + off_m * stride_dm
+    if DIVISIBLE_M:
+        tl.store(d_ptrs, delta)
+    else:
+        tl.store(d_ptrs, delta, mask=mask_m)
+
+
+@triton.jit
+def _bwd_kv_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale, DO,
+    DK, DV, DLOG_LAMBDA,
+    L,
+    D,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dkz, stride_dkh, stride_dkn, stride_dkk,
+    stride_dvz, stride_dvh, stride_dvn, stride_dvk,
+    stride_dlog_lambda_z, stride_dlog_lambda_h, stride_dlog_lambda_n,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_n = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+    log2e: tl.constexpr = 1.4426950408889634
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    DO += off_z * stride_doz + off_h * stride_doh
+
+    # offset pointers for batch/head
+    DK += off_z * stride_dkz + off_h * stride_dkh
+    DV += off_z * stride_dvz + off_h * stride_dvh
+    DLOG_LAMBDA += off_z * stride_dlog_lambda_z + off_h * stride_dlog_lambda_h
+
+    # offset pointers for batch/head
+    D += (off_z * H + off_h) * M
+    L += (off_z * H + off_h) * M
+
+    if CAUSAL:
+        lo = tl.maximum(start_n * BLOCK_N - P_SEQ, 0)
+        lo = (lo // BLOCK_M) * BLOCK_M
+    else:
+        lo = 0
+
+    offs_m_init = lo + tl.arange(0, BLOCK_M)
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_base = tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m_init[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m_init) * stride_log_lambda_n # (BLOCK_N, BLOCK_DMODEL)
+    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk) # (BLOCK_N, BLOCK_DMODEL)
+    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n * stride_log_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+    do_ptrs = DO + (offs_m_init[:, None] * stride_dom + offs_k[None, :] * stride_dok) # (BLOCK_M, BLOCK_DMODEL)
+
+    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_k[None, :] * stride_dvk) # (BLOCK_N, BLOCK_DMODEL)
+    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_k[None, :] * stride_dkk) # (BLOCK_N, BLOCK_DMODEL)
+    dlog_lambda_in_ptrs = DLOG_LAMBDA + (offs_n * stride_dlog_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+
+    # k and v stay in SRAM throughout
+    if DIVISIBLE_N:
+        v = tl.load(v_ptrs)
+        k = tl.load(k_ptrs)
+        log_lambda_in = tl.load(log_lambda_in_ptrs)
+    else:
+        mask_n = offs_n < N
+        v = tl.load(v_ptrs, mask=mask_n[:, None])
+        k = tl.load(k_ptrs, mask=mask_n[:, None])
+        log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n)
+
+    # If the N block doesn't contain seq_start, no need to loop
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        hi = tl.where(start_n * BLOCK_N + BLOCK_N >= seq_start - 1, M, lo)
+    else:
+        hi = M
+
+    # initialize dk amd dv
+    dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dlog_lambda_in = tl.zeros([BLOCK_N], dtype=tl.float32)
+
+    # loop over a col
+    for start_m in range(lo, hi, BLOCK_M):
+        start_m = tl.multiple_of(start_m, BLOCK_M)
+        offs_m = start_m + offs_m_base
+        causal_mask = (P_SEQ + offs_m[None, :]) >= (offs_n[:, None]) # (BLOCK_M, BLOCK_N)
+
+        # load q1, k1, q2, k2, v, do on-chip
+        if DIVISIBLE_M:
+            q = tl.load(q_ptrs)
+            log_lambda_out = tl.load(log_lambda_out_ptrs)
+        else:
+            mask_m = offs_m < M
+            valid_mask = mask_m[None, :] # & mask_n
+            q = tl.load(q_ptrs, mask=mask_m[:, None])
+            log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m)
+        # recompute p = softmax(qk * sm_scale, dim=-1)
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        sT = tl.dot(k, tl.trans(q), input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[None, :] - log_lambda_in[:, None]
+        sT += decay_bias * log2e
+        # NOTE: since softmax in backward is pointwise, the normalizer has been saved in fwd)
+        # So masking on s is not needed.
+        # s = tl.where(valid_mask, s , float("-inf"))
+        # if CAUSAL:
+        #     s = tl.where(causal_mask, s, float("-inf"))
+
+        # -- recompute p ---
+        if DIVISIBLE_M:
+            l = tl.load(L + offs_m)
+        else:
+            l = tl.load(L + offs_m, mask=mask_m)
+        pT = tl.math.exp2(sT - l[None, :] * log2e) # (BLOCK_M, BLOCK_N)
+
+        if not DIVISIBLE_M:
+            pT = tl.where(valid_mask, pT, 0.0)
+        if CAUSAL:
+            pT = tl.where(causal_mask, pT, 0.0)
+
+        # compute dv = dot(p, do)
+        if DIVISIBLE_M:
+            do = tl.load(do_ptrs)
+        else:
+            do = tl.load(do_ptrs, mask=mask_m[:, None]) # (BLOCK_M, BLOCK_DMODEL)
+
+
+        dv += tl.dot(pT.to(input_dtype), do, input_precision="ieee") # (BLOCK_N, BLOCK_DMODEL)  # still correct
+
+        # compute dp = dot(v, do)
+        if DIVISIBLE_M:
+            delta = tl.load(D + offs_m)
+        else:
+            delta = tl.load(D + offs_m, mask=mask_m)
+        # dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        dpT = tl.dot(v, tl.trans(do), input_precision="ieee")
+
+
+        # compute ds = p * (dp - delta[:, None])
+        dsT = pT * (dpT - delta[None, :]) # (BLOCK_M, BLOCK_N)
+
+        if not DIVISIBLE_M:
+            dsT = tl.where(valid_mask, dsT, 0.0)
+        if CAUSAL:
+            dsT = tl.where(causal_mask, dsT, 0.0)
+
+        # compute dk = dot(ds.T, q) masking
+        dk += tl.dot(dsT.to(input_dtype), q, input_precision="ieee")
+        dlog_lambda_in += -tl.sum(dsT, axis=1)
+
+        # increment pointers
+        q_ptrs += BLOCK_M * stride_qm
+        log_lambda_out_ptrs += BLOCK_M * stride_log_lambda_n
+        do_ptrs += BLOCK_M * stride_dom
+
+    dk *= sm_scale
+    if HAS_SEQ_START:
+        # Mask out 
+        seq_mask = (offs_n >= seq_start)
+        dk = tl.where(seq_mask[:, None], dk, 0.0)
+        dv = tl.where(seq_mask[:, None], dv, 0.0)
+        dlog_lambda_in = tl.where(seq_mask, dlog_lambda_in, 0.0)
+    if DIVISIBLE_N:
+        tl.store(dk_ptrs, dk.to(input_dtype)) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dv_ptrs, dv.to(input_dtype)) # (BLOCK_N, BLOCK_DMODEL,)
+        tl.store(dlog_lambda_in_ptrs, dlog_lambda_in.to(tl.float32)) # (BLOCK_N, BLOCK_DMODEL,)
+    else:
+        tl.store(dk_ptrs, dk.to(input_dtype), mask=mask_n[:, None]) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dv_ptrs, dv.to(input_dtype), mask=mask_n[:, None]) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dlog_lambda_in_ptrs, dlog_lambda_in.to(tl.float32), mask=mask_n) # (BLOCK_N, BLOCK_DMODEL,)
+
+
+@triton.jit
+def _bwd_q_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale, DO,
+    DQ, DLOG_LAMBDA,
+    L,
+    D,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dqz, stride_dqh, stride_dqm, stride_dqk,
+    stride_dlog_lambda_z, stride_dlog_lambda_h, stride_dlog_lambda_n,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr, LARGER_M: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_m = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    log2e: tl.constexpr = 1.4426950408889634
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    DO += off_z * stride_doz + off_h * stride_doh
+    D += (off_z * H + off_h) * M
+    L += (off_z * H + off_h) * M
+
+    # offset pointers for batch/head
+    DQ += off_z * stride_dqz + off_h * stride_dqh
+    DLOG_LAMBDA += off_z * stride_dlog_lambda_z + off_h * stride_dlog_lambda_h
+
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m) * stride_log_lambda_n
+
+    dq_ptrs = DQ + (offs_m[:, None] * stride_dqm + offs_k[None, :] * stride_dqk) # (BLOCK_M, BLOCK_DMODEL)
+    dlog_lambda_out_ptrs = DLOG_LAMBDA + (P_SEQ + offs_m) * stride_dlog_lambda_n
+    do_ptrs = DO + (offs_m[:, None] * stride_dom + offs_k[None, :] * stride_dok) # (BLOCK_M, BLOCK_DMODEL)
+
+    # pointer to row-wise quantities in value-like data
+    d_ptrs = D + offs_m
+    l_ptrs = L + offs_m
+
+    # load q: it will stay in SRAM throughout
+    if DIVISIBLE_M:
+        q = tl.load(q_ptrs)
+        do = tl.load(do_ptrs)
+        delta = tl.load(d_ptrs)
+        l = tl.load(l_ptrs)
+        log_lambda_out = tl.load(log_lambda_out_ptrs)
+    else:
+        mask_m = offs_m < M
+        q = tl.load(q_ptrs, mask=mask_m[:, None])
+        do = tl.load(do_ptrs, mask=mask_m[:, None])
+        delta = tl.load(d_ptrs, mask=mask_m)
+        l = tl.load(l_ptrs, mask=mask_m)
+        log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m)
+
+    # initialize dq
+    dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    dlog_lambda_out = tl.zeros([BLOCK_M], dtype=tl.float32)
+
+    # loop over k, v and update accumulator
+    # see note "Loop-Bound-For-N"
+    if CAUSAL:
+        hi = tl.minimum(N, P_SEQ + (start_m + 1) * BLOCK_M)
+        if LARGER_M:
+            hi = tl.maximum(0, hi)
+    else:
+        hi = N
+
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_n_init = offs_n_base
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        lo = tl.minimum(seq_start, hi)
+        lo = (lo // BLOCK_N) * BLOCK_N
+        offs_n_init += lo
+    else:
+        lo = 0
+    k_ptrs = K + (offs_n_init[:, None] * stride_kn + offs_k[None, :] * stride_kk) # (BLOCK_N, BLOCK_DMODEL)
+    v_ptrs = V + (offs_n_init[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n_init * stride_log_lambda_n)
+
+    # loop over a row
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + offs_n_base
+
+        # load k1, k2, v on chip
+        if DIVISIBLE_N:
+            v = tl.load(v_ptrs)
+            k = tl.load(k_ptrs)
+            log_lambda_in = tl.load(log_lambda_in_ptrs)
+        else:
+            mask_n = offs_n < N
+            v = tl.load(v_ptrs, mask=mask_n[:, None])
+            k = tl.load(k_ptrs, mask=mask_n[:, None])
+            log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n)
+
+
+        # recompute p = softmax(qk * sm_scale, dim=-1)
+        if not DIVISIBLE_N:
+            valid_mask = mask_n[None, :] # & mask_m[:, None]
+        if CAUSAL:
+            causal_mask = (P_SEQ + offs_m[:, None]) >= (offs_n[None, :]) # (BLOCK_M, BLOCK_N)
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        s = tl.dot(q, tl.trans(k), input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[:, None] - log_lambda_in[None, :]
+        s += decay_bias * log2e
+
+        # NOTE: since softmax in backward is pointwise, the normalizer has been saved in fwd)
+        # So masking on s is not needed.
+        # if CAUSAL:
+        #     s = tl.where(causal_mask & valid_mask, s, float("-inf"))
+        # else:
+        #     s = tl.where(valid_mask, s, float("-inf"))
+        p = tl.math.exp2(s - l[:, None] * log2e) # (BLOCK_M, BLOCK_N)
+
+        # compute dp = dot(v, do)
+        # dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        dp = tl.dot(do.to(input_dtype), tl.trans(v), input_precision="ieee")
+
+
+        # no need to mask dp
+        # if CAUSAL:
+        #     dp = tl.where(causal_mask & valid_mask, dp, 0.0)
+        # else:
+        #     dp = tl.where(valid_mask, dp, 0.0)
+
+        # compute ds = p * (dp - delta[:, None])
+        # move scale out to dq at last
+        ds = p * (dp - delta[:, None]) # (BLOCK_M, BLOCK_N)
+
+        # mask ds to ensure no small values
+        if not DIVISIBLE_N:
+            ds = tl.where(valid_mask, ds, 0.0)
+        if CAUSAL:
+            ds = tl.where(causal_mask, ds, 0.0)
+        if HAS_SEQ_START:
+            ds = tl.where(offs_n[None, :] >= seq_start, ds, 0.0)
+
+        dq += tl.dot(ds.to(input_dtype), k, input_precision="ieee")
+        dlog_lambda_out += tl.sum(ds, axis=1)
+
+        # increment pointers
+        k_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vn
+        log_lambda_in_ptrs += BLOCK_N * stride_log_lambda_n
+
+    dq *= sm_scale
+    if DIVISIBLE_M:
+        tmp = tl.load(dlog_lambda_out_ptrs)
+    else:
+        tmp = tl.load(dlog_lambda_out_ptrs, mask=mask_m)
+    dlog_lambda_out += tmp
+    if DIVISIBLE_M:
+        tl.store(dq_ptrs, dq.to(input_dtype))
+        tl.store(dlog_lambda_out_ptrs, dlog_lambda_out)
+    else:
+        tl.store(dq_ptrs, dq.to(input_dtype), mask=mask_m[:, None])
+        tl.store(dlog_lambda_out_ptrs, dlog_lambda_out, mask=mask_m)
+
+
+
+@pytest.mark.parametrize("Z, H, M, N, HEAD_DIM", [(4, 2, 1020, 2098, 64), (4, 2, 1024, 2048, 64)])
+@pytest.mark.parametrize("causal", [True])
+def test_op(Z, H, M, N, HEAD_DIM, causal, dtype=torch.bfloat16):
+    torch.manual_seed(24)
+    q = (torch.empty((Z, H, M, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    k = (torch.empty((Z, H, N, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    v = (torch.empty((Z, H, N, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    fgate_logit = torch.empty((Z, H, N), dtype=torch.float32, device="cuda").uniform_(5, 10)
+    log_fgate = torch.nn.functional.logsigmoid(fgate_logit).requires_grad_()
+    seq_start = torch.randint(low=0, high=N, size=(Z,), dtype=torch.long, device="cuda")
+    # seq_start = torch.randint(low=0, high=10, size=(Z,), dtype=torch.long, device="cuda")
+    # seq_start = torch.full(fill_value=0, size=(Z,), dtype=torch.long, device="cuda")
+    sm_scale = 0.5
+    dout = torch.randn_like(q)
+    # reference implementation
+    P_SEQ = N - M
+    mask = torch.tril(torch.ones((M, N), device="cuda"), diagonal=P_SEQ)
+    p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
+    p = p.float()
+
+    log_lambda = torch.cumsum(log_fgate, dim=-1)
+    decay_bias = log_lambda[..., -M:, None] - log_lambda[..., None, :]
+    p = p + decay_bias
+    if causal:
+        p[:, :, mask == 0] = float("-inf")
+
+    attention_mask = torch.arange(N, device="cuda") < seq_start[:, None, None, None]
+    p = torch.where(attention_mask, float("-inf"), p)
+    p = torch.softmax(p.float(), dim=-1).to(dtype)
+    p = p.clone()
+    p[torch.isnan(p)] = 0.0
+    # p = torch.exp(p)
+    ref_out = torch.matmul(p, v)
+    ref_out.backward(dout)
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dlog_fgate, log_fgate.grad = log_fgate.grad.clone(), None
+    # triton implementation
+    tri_out = forgetting_attention(q, k, v, log_fgate, head_first=True, seq_start=seq_start, sm_scale=sm_scale)
+    tri_out = tri_out.to(dtype)
+
+    tri_out.backward(dout)
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dlog_fgate, log_fgate.grad = log_fgate.grad.clone(), None
+    # compare
+    # assert torch.allclose(tri_log_normalizer[~torch.isnan(tri_log_normalizer)], ref_log_normalizer[~torch.isnan(ref_log_normalizer)], atol=1e-2, rtol=0)
+    assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0), (ref_out - tri_out).abs().max()
+    rtol = 0
+    # Relative tolerance workaround for known hardware limitation of MI200 GPU.
+    # For details see https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
+    # if torch.version.hip is not None and triton.runtime.driver.active.get_current_target().arch == "gfx90a":
+        # rtol = 1e-2
+    assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=rtol), (ref_dv - tri_dv).abs().max()
+    assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=rtol), (ref_dk - tri_dk).abs().max()
+    assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=rtol), (ref_dq - tri_dq).abs().max()
+    assert torch.allclose(ref_dlog_fgate, tri_dlog_fgate, atol=1e-2, rtol=rtol), (ref_dlog_fgate - tri_dlog_fgate).abs().max()
+
+try:
+    from flash_attn.flash_attn_interface import \
+        flash_attn_qkvpacked_func as flash_attn_func
+    HAS_FLASH = True
+except BaseException:
+    HAS_FLASH = False
+
+TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
+BATCH, N_HEADS, HEAD_DIM = 4, 32, 128
+# vary seq length for fixed head and batch=4
+configs = []
+for mode in ["fwd", "bwd"]:
+# for mode in ["bwd"]:
+    # for causal in [True, False]:
+    for causal in [True]:
+        if mode == "bwd" and not causal:
+            continue
+        configs.append(
+            triton.testing.Benchmark(
+                x_names=["N_CTX"],
+                # x_vals=[2**i for i in range(10, 15)],
+                x_vals=[2**i for i in range(14, 15)],
+                line_arg="provider",
+                # line_vals=["triton-fp16", "flag"] + (["flash"] if HAS_FLASH else []),
+                # line_names=["Triton [FP16]", "Flag"] + (["Flash-2"] if HAS_FLASH else []),
+                line_vals=["flag"] + (["flash"] if HAS_FLASH else []),
+                line_names=["Flag"] + (["Flash-2"] if HAS_FLASH else []),
+                styles=[("red", "-"), ("blue", "-"), ("green", "-")],
+                ylabel="ms",
+                plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}",
+                args={
+                    "H": N_HEADS,
+                    "BATCH": BATCH,
+                    "HEAD_DIM": HEAD_DIM,
+                    "mode": mode,
+                    "causal": causal,
+                },
+            ))
+
+
+@triton.testing.perf_report(configs)
+def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, causal, mode, provider, device="cuda"):
+    assert mode in ["fwd", "bwd"]
+    warmup = 25
+    rep = 100
+    dtype = torch.bfloat16
+    if "flag" in provider:
+        q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fgate_logit = torch.empty((BATCH, H, N_CTX), dtype=torch.float32, device="cuda").uniform_(5, 10)
+        log_fgate = torch.nn.functional.logsigmoid(fgate_logit).requires_grad_()
+        # if mode == "fwd" and "fp8" in provider:
+        #     q = q.to(torch.float8_e5m2)
+        #     k = k.to(torch.float8_e5m2)
+        #     v = v.permute(0, 1, 3, 2).contiguous()
+        #     v = v.permute(0, 1, 3, 2)
+        #     v = v.to(torch.float8_e5m2)
+        sm_scale = 1.3
+        fn = lambda: forgetting_attention(q, k, v, log_fgate, head_first=True, sm_scale=sm_scale)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    if provider == "flash":
+        qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fn = lambda: flash_attn_func(qkv, causal=causal)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM
+    total_flops = 2 * flops_per_matmul
+    if causal:
+        total_flops *= 0.5
+    if mode == "bwd":
+        total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
+    return total_flops / ms * 1e-9
+
+
+if __name__ == "__main__":
+    # only works on post-Ampere GPUs right now
+    bench_flash_attention.run(save_path=".", print_data=True)

ops/.ipynb_checkpoints/forgetting_attention_std-checkpoint.py

@@ -0,0 +1,72 @@
+"""
+Forgetting Attention - 标准 Softmax 版本
+在 forgetting_attention.py 最后添加这个函数
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def forgetting_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_fgate: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """标准 Softmax 版本的 Forgetting Attention"""
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+        log_fgate = rearrange(log_fgate, "b t h -> b h t")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # 计算 QK 分数
+    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # 处理 seq_start
+    log_fgate_masked = log_fgate.float()
+    if seq_start is not None:
+        log_fgate_masked = log_fgate_masked.clone()
+        mask_idx = torch.arange(T_k, device=q.device)[None, None, :] < seq_start[:, None, None]
+        log_fgate_masked[mask_idx] = 0.0
+    
+    # 计算累积衰减
+    log_lambda = torch.cumsum(log_fgate_masked, dim=-1)
+    decay_bias = log_lambda[:, :, :T_q, None] - log_lambda[:, :, None, :]
+    scores = scores + decay_bias
+    
+    # Causal mask
+    P_SEQ = T_k - T_q
+    causal_mask = torch.triu(torch.ones((T_q, T_k), dtype=torch.bool, device=q.device), diagonal=P_SEQ + 1)
+    scores = scores.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # seq_start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        scores = scores.masked_fill(seq_mask, float('-inf'))
+    
+    # Softmax
+    attn = F.softmax(scores, dim=-1)
+    attn = torch.nan_to_num(attn, 0.0)
+    
+    # 计算输出
+    out = torch.matmul(attn.to(v.dtype), v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out

ops/.ipynb_checkpoints/geometric_attention_std-checkpoint.py

@@ -0,0 +1,179 @@
+"""
+Geometric Attention - 标准 Softmax 版本
+基于论文 "The Neural Data Router" (Csordás et al., 2022)
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def geometric_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    标准 Softmax 版本的 Geometric Attention
+    
+    Args:
+        q: Query tensor [B, T, H, D] or [B, H, T, D] if head_first
+        k: Key tensor [B, T, H, D] or [B, H, T, D] if head_first
+        v: Value tensor [B, T, H, D] or [B, H, T, D] if head_first
+        head_first: 是否head维度在前
+        seq_start: 序列起始位置 [B]
+        sm_scale: scaling factor，默认 1/sqrt(D)
+        normalize: 是否归一化attention weights
+    
+    Returns:
+        output: [B, T, H, D] or [B, H, T, D] if head_first
+    """
+    
+    # Rearrange to head_first format
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Step 1: 计算 content-based logits
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    # logits: [B, H, T_q, T_k]
+    
+    # Step 2: Mask diagonal (不允许attend到自己)
+    if T_q == T_k:
+        diag_mask = torch.eye(T_q, dtype=torch.bool, device=q.device)
+        logits = logits.masked_fill(diag_mask[None, None, :, :], float('-inf'))
+    
+    # Step 3: 处理 seq_start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        logits = logits.masked_fill(seq_mask, float('-inf'))
+    
+    # Step 4: Causal mask (如果需要)
+    # 注意：geometric attention论文中没有causal，如果你的任务需要可以取消注释
+    # P_SEQ = T_k - T_q
+    # causal_mask = torch.triu(torch.ones((T_q, T_k), dtype=torch.bool, device=q.device), diagonal=P_SEQ + 1)
+    # logits = logits.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # Step 5: Geometric weighting (核心算法)
+    attn_weights = geometric_weighting(logits, normalize=normalize)
+    
+    # Step 6: 应用attention到values
+    out = torch.matmul(attn_weights.to(v.dtype), v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out
+
+
+def geometric_weighting(
+    logits: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    计算geometric attention weights
+    
+    实现论文中的 Equation 7:
+    A[i,j] = P[i,j] * ∏(1 - P[i,k]) for k closer to i than j
+    
+    Args:
+        logits: [B, H, T_q, T_k] attention logits
+        normalize: 是否归一化
+    
+    Returns:
+        weights: [B, H, T_q, T_k] attention weights
+    """
+    B, H, T_q, T_k = logits.shape
+    
+    # Step 1: Sigmoid to get matching probabilities
+    P = torch.sigmoid(logits)  # [B, H, T_q, T_k]
+    
+    # Step 2: 使用 log-space 计算（数值稳定）
+    log_P = torch.log(P + 1e-10)
+    log_one_minus_P = torch.log(1.0 - P + 1e-10)
+    
+    # Step 3: 简化版本 - 使用cumsum实现几何分布
+    # 这是一个高效的近似，避免了显式的循环
+    
+    # 对于每个位置i，计算其左侧所有位置的log(1-P)累积和
+    log_decay_left = log_one_minus_P.cumsum(dim=-1)
+    
+    # 计算weights（简化版）
+    # 完整版本需要根据距离动态选择区间，这里用一个高效近似
+    weights = torch.exp(log_P + log_decay_left.roll(1, dims=-1))
+    
+    # 第一个位置特殊处理（没有左侧元素）
+    # 避免inplace操作
+    weights_first = P[:, :, :, :1]  # 获取第一列
+    weights = torch.cat([weights_first, weights[:, :, :, 1:]], dim=-1)
+    
+    # Step 4: 归一化（可选）
+    if normalize:
+        weights = F.normalize(weights, p=1, dim=-1)
+    
+    # 处理NaN（如果所有位置都是-inf）
+    weights = torch.nan_to_num(weights, 0.0)
+    
+    return weights
+
+
+def geometric_weighting_full(
+    logits: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    完整版geometric weighting（更慢但更准确）
+    
+    仅在需要最高精度时使用，训练时建议用上面的简化版
+    """
+    B, H, T_q, T_k = logits.shape
+    device = logits.device
+    
+    P = torch.sigmoid(logits)
+    log_P = torch.log(P + 1e-10)
+    log_one_minus_P = torch.log(1.0 - P + 1e-10)
+    
+    # 初始化weights
+    weights = torch.zeros_like(P)
+    
+    # 对每个(i,j)计算geometric weight
+    for i in range(T_q):
+        for j in range(T_k):
+            # 找出比j更接近i的所有位���k
+            if i < j:
+                # 向右看：closer positions are [i+1, ..., j-1]
+                closer_positions = range(i + 1, j)
+            elif i > j:
+                # 向左看：closer positions are [j+1, ..., i-1]
+                closer_positions = range(j + 1, i)
+            else:
+                # i == j (对角线)，已经在外面mask掉了
+                continue
+            
+            # 计算 ∏(1 - P[i,k]) in log-space
+            log_prod = sum(log_one_minus_P[:, :, i, k] for k in closer_positions) if closer_positions else 0.0
+            
+            # weights[i,j] = P[i,j] * ∏(1 - P[i,k])
+            weights[:, :, i, j] = torch.exp(log_P[:, :, i, j] + log_prod)
+    
+    if normalize:
+        weights = F.normalize(weights, p=1, dim=-1)
+    
+    weights = torch.nan_to_num(weights, 0.0)
+    
+    return weights 

ops/.ipynb_checkpoints/sliding_window_attention_std-checkpoint.py

@@ -0,0 +1,88 @@
+"""
+Sliding Window / Hard Attention
+Based on "Context Limitations Make Neural Language Models More Human-Like"
+(Kuribayashi et al., 2022)
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def sliding_window_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    window_size: int = 2,  # 默认2-gram（看前1个token）
+) -> torch.Tensor:
+    """
+    Sliding Window Attention
+    
+    硬截断：只能attend到最近window_size个token
+    """
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Compute logits
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # Create sliding window mask
+    mask = create_sliding_window_mask(T_q, T_k, window_size, device=q.device)
+    logits = logits.masked_fill(~mask, float('-inf'))
+    
+    # Seq start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        logits = logits.masked_fill(seq_mask, float('-inf'))
+    
+    # Standard softmax
+    weights = F.softmax(logits, dim=-1)
+    
+    # Apply to values
+    out = torch.matmul(weights, v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out
+
+
+def create_sliding_window_mask(
+    T_q: int,
+    T_k: int,
+    window_size: int,
+    device: torch.device
+) -> torch.Tensor:
+    """
+    创建sliding window mask
+    
+    window_size=1: 只看前1个token (2-gram)
+    window_size=2: 只看前2个token (3-gram)
+    """
+    # 基础causal mask
+    mask = torch.tril(torch.ones(T_q, T_k, dtype=torch.bool, device=device))
+    
+    # 应用window限制
+    if window_size > 0 and window_size < T_k:
+        for i in range(T_q):
+            # 只保留 [i-window_size+1, i] 范围
+            start = max(0, i - window_size + 1)
+            if start > 0:
+                mask[i, :start] = False
+    
+    return mask[None, None, :, :]  # [1, 1, T_q, T_k]

ops/.ipynb_checkpoints/stickbreaking_attention_std-checkpoint.py

@@ -0,0 +1,117 @@
+"""
+Stick-breaking Attention - ICLR 2025
+基于论文 "Scaling Stick-Breaking Attention" (Tan et al., 2025)
+简化的PyTorch实现（不使用Triton）
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def stickbreaking_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+    attend_current: bool = False,
+) -> torch.Tensor:
+    """
+    Stick-breaking attention
+    
+    Based on ICLR 2025 paper, simplified PyTorch implementation
+    A_{i,j} = exp(z_{i,j} - ∑_{k=i}^{j-1} softplus(z_{k,j}))
+    
+    Args:
+        q: query [B, T, H, D] or [B, H, T, D] if head_first
+        k: key [B, T, H, D] or [B, H, T, D] if head_first
+        v: value [B, T, H, D] or [B, H, T, D] if head_first
+        attend_current: whether to attend to current position
+        normalize: whether to normalize attention weights
+    """
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Compute logits: QK^T / sqrt(d)
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    # [B, H, T_q, T_k]
+    
+    # Causal mask (optional: mask diagonal if not attend_current)
+    if T_q == T_k and not attend_current:
+        diag_mask = torch.eye(T_q, dtype=torch.bool, device=q.device)
+        logits = logits.masked_fill(diag_mask[None, None, :, :], float('-inf'))
+    
+    # Seq start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        logits = logits.masked_fill(seq_mask, float('-inf'))
+    
+    # Stick-breaking weighting
+    attn_weights = stickbreaking_weighting(logits, normalize=normalize)
+    
+    # Apply attention to values
+    out = torch.matmul(attn_weights.to(v.dtype), v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out
+
+
+def stickbreaking_weighting(
+    logits: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    Compute stick-breaking attention weights
+    
+    From paper Equation 4:
+    A_{i,j} = exp(z_{i,j} - ∑_{k=i}^{j-1} log(1 + exp(z_{k,j})))
+    
+    Where log(1 + exp(x)) is softplus(x)
+    """
+    B, H, T_q, T_k = logits.shape
+    device = logits.device
+    
+    # Softplus: log(1 + exp(x))
+    # Numerically stable version from paper (Equation 5)
+    def softplus_stable(x):
+        # softplus(x) = log(1 + exp(x))
+        # When x > 15, exp(x) is huge, just return x
+        return torch.where(
+            x > 15.0,
+            x,
+            torch.log1p(torch.exp(torch.clamp(x, max=15.0)))
+        )
+    
+    # Compute softplus for all logits
+    logits_sp = softplus_stable(logits)  # [B, H, T_q, T_k]
+    
+    # For each query position, compute cumulative sum
+    # We need to accumulate from left to right (position i to j-1)
+    log_weights = torch.zeros_like(logits)
+    
+    for i in range(T_q):
+        # For query i, we compute attention to all keys j
+        z_i = logits[:, :, i, :]      # [B, H, T_k]
+        z_sp_i = logits_sp[:, :, i, :] # [B, H, T_k]
+        
+        # Cumulative sum of softplus
+        # csum[j] = ∑_{k=0}^{j} softplus(z_{i,k})
+        csum = z_sp_i.cumsum(dim=-1)

ops/.ipynb_checkpoints/vanilla_attention_std-checkpoint.py

@@ -0,0 +1,171 @@
+"""
+Vanilla Transformer 的标准 Softmax Attention
+用于替换 flash_attn 的实现
+"""
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional, Tuple
+
+def vanilla_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor, 
+    v: torch.Tensor,
+    causal: bool = True,
+    window_size: Optional[Tuple[int, int]] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """
+    标准 Softmax Attention，兼容 flash_attn_func 的输入格式
+    
+    Args:
+        q, k, v: [batch, seq_len, num_heads, head_dim] 格式
+        causal: 是否使用因果mask
+        window_size: 滑动窗口大小 (left, right)，(-1, -1) 表示无限制
+        sm_scale: softmax 缩放因子
+    
+    Returns:
+        output: [batch, seq_len, num_heads, head_dim] 格式
+    """
+    B, T_q, H, D = q.shape
+    T_k = k.shape[1]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # 转换为 [B, H, T, D] 格式进行计算
+    q = rearrange(q, 'b t h d -> b h t d')
+    k = rearrange(k, 'b t h d -> b h t d')
+    v = rearrange(v, 'b t h d -> b h t d')
+    
+    # 计算 attention scores
+    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # Causal mask
+    if causal:
+        P_SEQ = T_k - T_q  # 处理 KV cache 的情况
+        causal_mask = torch.triu(
+            torch.ones((T_q, T_k), dtype=torch.bool, device=q.device),
+            diagonal=P_SEQ + 1
+        )
+        scores = scores.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # Window mask (sliding window attention)
+    if window_size is not None and window_size != (-1, -1):
+        left_window, right_window = window_size
+        window_mask = torch.ones((T_q, T_k), dtype=torch.bool, device=q.device)
+        for i in range(T_q):
+            # 计算每个查询位置的有效窗口范围
+            start = max(0, i - left_window)
+            end = min(T_k, i + right_window + 1)
+            window_mask[i, start:end] = False
+        scores = scores.masked_fill(window_mask[None, None, :, :], float('-inf'))
+    
+    # Softmax
+    attn_weights = F.softmax(scores, dim=-1)
+    attn_weights = torch.nan_to_num(attn_weights, 0.0)
+    
+    # Apply attention to values
+    output = torch.matmul(attn_weights.to(v.dtype), v)
+    
+    # 转换回 [B, T, H, D] 格式
+    output = rearrange(output, 'b h t d -> b t h d')
+    
+    return output
+
+
+def vanilla_attention_varlen_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    causal: bool = True,
+    window_size: Optional[Tuple[int, int]] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """
+    变长序列的标准 Softmax Attention，兼容 flash_attn_varlen_func
+    
+    Args:
+        q: [total_q_tokens, num_heads, head_dim]
+        k: [total_k_tokens, num_kv_heads, head_dim]
+        v: [total_k_tokens, num_kv_heads, head_dim]
+        cu_seqlens_q: 累积序列长度 [batch_size + 1]
+        cu_seqlens_k: 累积序列长度 [batch_size + 1]
+        max_seqlen_q: 最大查询序列长度
+        max_seqlen_k: 最大键值序列长度
+    
+    Returns:
+        output: [total_q_tokens, num_heads, head_dim]
+    """
+    batch_size = cu_seqlens_q.shape[0] - 1
+    H = q.shape[1]
+    D = q.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    outputs = []
+    
+    # 逐批次处理
+    for b in range(batch_size):
+        q_start, q_end = cu_seqlens_q[b].item(), cu_seqlens_q[b+1].item()
+        k_start, k_end = cu_seqlens_k[b].item(), cu_seqlens_k[b+1].item()
+        
+        if q_start == q_end:  # 空序列
+            continue
+            
+        # 提取当前批次的 q, k, v
+        q_b = q[q_start:q_end]  # [T_q, H, D]
+        k_b = k[k_start:k_end]  # [T_k, H, D] 
+        v_b = v[k_start:k_end]  # [T_k, H, D]
+        
+        T_q = q_b.shape[0]
+        T_k = k_b.shape[0]
+        
+        # 转换为 [H, T, D] 格式
+        q_b = rearrange(q_b, 't h d -> h t d')
+        k_b = rearrange(k_b, 't h d -> h t d')
+        v_b = rearrange(v_b, 't h d -> h t d')
+        
+        # 计算 attention scores
+        scores = torch.matmul(q_b.float(), k_b.float().transpose(-2, -1)) * sm_scale
+        
+        # Causal mask
+        if causal:
+            P_SEQ = T_k - T_q
+            causal_mask = torch.triu(
+                torch.ones((T_q, T_k), dtype=torch.bool, device=q.device),
+                diagonal=P_SEQ + 1
+            )
+            scores = scores.masked_fill(causal_mask[None, :, :], float('-inf'))
+        
+        # Window mask
+        if window_size is not None and window_size != (-1, -1):
+            left_window, right_window = window_size
+            window_mask = torch.ones((T_q, T_k), dtype=torch.bool, device=q.device)
+            for i in range(T_q):
+                start = max(0, i - left_window)
+                end = min(T_k, i + right_window + 1)
+                window_mask[i, start:end] = False
+            scores = scores.masked_fill(window_mask[None, :, :], float('-inf'))
+        
+        # Softmax
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = torch.nan_to_num(attn_weights, 0.0)
+        
+        # Apply attention
+        output_b = torch.matmul(attn_weights.to(v_b.dtype), v_b)
+        
+        # 转换回 [T, H, D] 格式
+        output_b = rearrange(output_b, 'h t d -> t h d')
+        outputs.append(output_b)
+    
+    # 拼接所有批次的输出
+    output = torch.cat(outputs, dim=0)
+    
+    return output

ops/__init__.py

@@ -0,0 +1,3 @@
+
+# Framework mock for ndr compatibility
+from . import framework_mock

ops/direction_sensitive_geometric.py

@@ -0,0 +1,115 @@
+import torch
+from forgetting_transformer.ops.multi_head_attention import AttentionMask, MultiHeadAttentionBase, AttentionMergeMixin
+from typing import Optional
+from forgetting_transformer.ops.geometric_attention import geometric_attention_activation
+import math
+from forgetting_transformer.ops.multi_head_relative_pos_attention import FixedRelativeMultiheadAttentionBase, shift
+
+
+class DirectionSensitiveGeometricAttention(AttentionMergeMixin, FixedRelativeMultiheadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, global_pos_bias: bool = True,
+                 global_content_bias: bool = True, input_size: Optional[int] = None,
+                 output_size: Optional[int] = None, normalize_score: bool = True):
+        super(AttentionMergeMixin, self).__init__(state_size, n_heads, dropout, input_size)
+
+        self.data_to_kv = torch.nn.Linear(state_size, 2 * n_heads * self.projection_size, bias=False)
+        self.data_to_q = torch.nn.Linear(self.input_size, n_heads * self.projection_size, bias=False)
+        self.data_to_qp = torch.nn.Linear(self.input_size, n_heads * 2)
+
+        self.global_content_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                                   if global_content_bias else None
+
+        self.s_bias = torch.nn.Parameter(torch.full([1], 0.0))
+        self.scale = torch.nn.Parameter(torch.full([1], 1.0 / math.sqrt(self.projection_size)))
+        self.scale_pos = torch.nn.Parameter(torch.full([1], 1.0))
+        self.normalize_score = normalize_score
+
+        self.input_size = state_size if input_size is None else input_size
+
+        print(f"DirectionSensitiveGeometricAttention: normalize score: {normalize_score}")
+
+        super(DirectionSensitiveGeometricAttention, self).__init__(output_size)
+        self.reset_parameters()
+
+    def get_attention_scores(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor,
+                   pos_offset: int) -> torch.Tensor:
+
+        # content-content addressing
+        logits = torch.bmm(q_content, self.dropout(k_content).transpose(1, 2))
+
+        # directionality. Do scaling here, less flops.
+        prefer_back, prefer_front = (q_pos * self.scale_pos).unsqueeze(-2).expand(-1,-1,logits.shape[-1],-1).unbind(-1)
+        fpos = prefer_front.triu(1 + pos_offset) + prefer_back.tril(-1 + pos_offset)
+
+        logits = logits * self.scale + fpos + self.s_bias
+
+        logits = self.apply_logit_masks(logits.view(logits.shape[0] // self.n_heads, self.n_heads, *logits.shape[1:]), mask).flatten(0,1)
+
+        logits.masked_fill_(torch.eye(logits.shape[-1], device=logits.device, dtype=torch.bool)[pos_offset : pos_offset + logits.shape[-2]], float("-inf"))
+
+        return geometric_attention_activation(logits, mask, pos_offset, normalize=self.normalize_score)
+
+    def add_head_specific_bias(self, data: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor:
+        # data [batch * n_heads, len, c]
+        # bias [n_heads, c]
+        return (data.view(-1, bias.shape[0], *data.shape[1:]) + bias.unsqueeze(1).type_as(data)).view_as(data) \
+               if bias is not None else data
+
+    def _attention(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor,
+                   v: torch.Tensor, pos_offset: int) -> [torch.Tensor, torch.Tensor]:
+
+        scores = self.get_attention_scores(mask, q_content, k_content, q_pos, pos_offset)
+
+        # Scores shape: [n_batch * n_heads, n_out, n_in]
+        return self._attention_read(mask, scores, v)
+
+    def forward(self, curr_state: torch.Tensor, attend_to: torch.Tensor, mask: Optional[AttentionMask],
+                pos_offset: int = 0, need_weights: bool = False):
+        # curr_state: [batch_size, out_len, c]
+        # attend_to: [batch_size, in_len, c]
+        batch_size, in_len = attend_to.shape[0:2]
+        out_len = curr_state.shape[1]
+
+        k_content, v = self.transform_data(attend_to, self.data_to_kv, 2)
+        q, = self.transform_data(curr_state, self.data_to_q, 1)
+        q_pos, = self.transform_data(curr_state, self.data_to_qp, 1)
+
+        q_content = self.add_head_specific_bias(q, self.global_content_bias)
+
+        data, scores = self.merged_attention(batch_size, out_len, mask, q_content, k_content, q_pos, v,
+                                             pos_offset, need_weights=need_weights)
+
+        if need_weights:
+            return data, scores
+        else:
+            return data
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.data_to_q.weight)
+        torch.nn.init.xavier_uniform_(self.pos_to_pq.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight[:self.projection_size * self.n_heads])
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight[self.projection_size * self.n_heads:])
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)
+
+
+class DirectionSensitiveGeometricAttentionMyInit(DirectionSensitiveGeometricAttention):
+    def xavier_manual_(self, tensor: torch.Tensor, fan_in: int, fan_out: int, gain: float = 1) -> torch.Tensor:
+        std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
+        a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+
+        return torch.nn.init._no_grad_uniform_(tensor, -a, a)
+
+    def reset_parameters(self):
+        self.xavier_manual_(self.data_to_q.weight, self.state_size, self.projection_size)
+        self.xavier_manual_(self.pos_to_pq.weight, self.state_size, 2)
+        self.xavier_manual_(self.data_to_kv.weight, self.state_size, self.projection_size)
+        self.xavier_manual_(self.multi_head_merge.weight, self.projection_size, self.state_size)
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)

ops/direction_sensitive_geometric.py.bak

@@ -0,0 +1,115 @@
+import torch
+from .multi_head_attention import AttentionMask, MultiHeadAttentionBase, AttentionMergeMixin
+from typing import Optional
+from .geometric_attention import geometric_attention_activation
+import math
+from .multi_head_relative_pos_attention import FixedRelativeMultiheadAttentionBase, shift
+
+
+class DirectionSensitiveGeometricAttention(AttentionMergeMixin, FixedRelativeMultiheadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, global_pos_bias: bool = True,
+                 global_content_bias: bool = True, input_size: Optional[int] = None,
+                 output_size: Optional[int] = None, normalize_score: bool = True):
+        super(AttentionMergeMixin, self).__init__(state_size, n_heads, dropout, input_size)
+
+        self.data_to_kv = torch.nn.Linear(state_size, 2 * n_heads * self.projection_size, bias=False)
+        self.data_to_q = torch.nn.Linear(self.input_size, n_heads * self.projection_size, bias=False)
+        self.data_to_qp = torch.nn.Linear(self.input_size, n_heads * 2)
+
+        self.global_content_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                                   if global_content_bias else None
+
+        self.s_bias = torch.nn.Parameter(torch.full([1], 0.0))
+        self.scale = torch.nn.Parameter(torch.full([1], 1.0 / math.sqrt(self.projection_size)))
+        self.scale_pos = torch.nn.Parameter(torch.full([1], 1.0))
+        self.normalize_score = normalize_score
+
+        self.input_size = state_size if input_size is None else input_size
+
+        print(f"DirectionSensitiveGeometricAttention: normalize score: {normalize_score}")
+
+        super(DirectionSensitiveGeometricAttention, self).__init__(output_size)
+        self.reset_parameters()
+
+    def get_attention_scores(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor,
+                   pos_offset: int) -> torch.Tensor:
+
+        # content-content addressing
+        logits = torch.bmm(q_content, self.dropout(k_content).transpose(1, 2))
+
+        # directionality. Do scaling here, less flops.
+        prefer_back, prefer_front = (q_pos * self.scale_pos).unsqueeze(-2).expand(-1,-1,logits.shape[-1],-1).unbind(-1)
+        fpos = prefer_front.triu(1 + pos_offset) + prefer_back.tril(-1 + pos_offset)
+
+        logits = logits * self.scale + fpos + self.s_bias
+
+        logits = self.apply_logit_masks(logits.view(logits.shape[0] // self.n_heads, self.n_heads, *logits.shape[1:]), mask).flatten(0,1)
+
+        logits.masked_fill_(torch.eye(logits.shape[-1], device=logits.device, dtype=torch.bool)[pos_offset : pos_offset + logits.shape[-2]], float("-inf"))
+
+        return geometric_attention_activation(logits, mask, pos_offset, normalize=self.normalize_score)
+
+    def add_head_specific_bias(self, data: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor:
+        # data [batch * n_heads, len, c]
+        # bias [n_heads, c]
+        return (data.view(-1, bias.shape[0], *data.shape[1:]) + bias.unsqueeze(1).type_as(data)).view_as(data) \
+               if bias is not None else data
+
+    def _attention(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor,
+                   v: torch.Tensor, pos_offset: int) -> [torch.Tensor, torch.Tensor]:
+
+        scores = self.get_attention_scores(mask, q_content, k_content, q_pos, pos_offset)
+
+        # Scores shape: [n_batch * n_heads, n_out, n_in]
+        return self._attention_read(mask, scores, v)
+
+    def forward(self, curr_state: torch.Tensor, attend_to: torch.Tensor, mask: Optional[AttentionMask],
+                pos_offset: int = 0, need_weights: bool = False):
+        # curr_state: [batch_size, out_len, c]
+        # attend_to: [batch_size, in_len, c]
+        batch_size, in_len = attend_to.shape[0:2]
+        out_len = curr_state.shape[1]
+
+        k_content, v = self.transform_data(attend_to, self.data_to_kv, 2)
+        q, = self.transform_data(curr_state, self.data_to_q, 1)
+        q_pos, = self.transform_data(curr_state, self.data_to_qp, 1)
+
+        q_content = self.add_head_specific_bias(q, self.global_content_bias)
+
+        data, scores = self.merged_attention(batch_size, out_len, mask, q_content, k_content, q_pos, v,
+                                             pos_offset, need_weights=need_weights)
+
+        if need_weights:
+            return data, scores
+        else:
+            return data
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.data_to_q.weight)
+        torch.nn.init.xavier_uniform_(self.pos_to_pq.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight[:self.projection_size * self.n_heads])
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight[self.projection_size * self.n_heads:])
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)
+
+
+class DirectionSensitiveGeometricAttentionMyInit(DirectionSensitiveGeometricAttention):
+    def xavier_manual_(self, tensor: torch.Tensor, fan_in: int, fan_out: int, gain: float = 1) -> torch.Tensor:
+        std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
+        a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+
+        return torch.nn.init._no_grad_uniform_(tensor, -a, a)
+
+    def reset_parameters(self):
+        self.xavier_manual_(self.data_to_q.weight, self.state_size, self.projection_size)
+        self.xavier_manual_(self.pos_to_pq.weight, self.state_size, 2)
+        self.xavier_manual_(self.data_to_kv.weight, self.state_size, self.projection_size)
+        self.xavier_manual_(self.multi_head_merge.weight, self.projection_size, self.state_size)
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)

ops/forgetting_attention.py

@@ -0,0 +1,1138 @@
+"""
+Implementation of Forgetting Attention.
+
+Our code is adapted from https://github.com/FlagOpen/FlagAttention/blob/ee91638dec6da8c00c4113d179f469e0ffcd5852/src/flag_attn/flash.py. The code is modified to implement Forgetting Attention.
+
+The original license info from FlagAttention:
+
+Copyright 2023 BAAI
+
+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 pytest
+import math
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+from typing import Optional
+
+
+__all__ = ["forgetting_attention"]
+
+
+# File flash.py
+def maybe_contiguous(x):
+    # only when the inner most dimension is contiguous can LDGSTS be used
+    # so inner-dimension contiguity is enforced.
+    return x.contiguous() if x.stride(-1) != 1 else x
+
+def rounded_multiple(a, b):
+    return (a + b - 1) // b * b
+
+# --------------------------- public API ---------------------------
+class ForgettingAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, log_fgate, seq_start, causal, sm_scale, return_log_normalizer):
+        assert causal, "Only causal attention is supported"
+        Dq, Dk, Dv = q.shape[-1], k.shape[-1], v.shape[-1]
+        assert Dq == Dk == Dv, "feature size of q, k, v should be equal"
+        assert Dk in {16, 32, 64, 128}, "We only support head dims in {16, 32, 64, 128}"
+
+        B, H, M, D = q.shape
+        if seq_start is not None:
+            has_seq_start = True
+            assert seq_start.shape == (B,)
+        else:
+            has_seq_start = False
+            seq_start = torch.zeros((B,), device=q.device, dtype=torch.long)
+        N = k.shape[2]
+        assert log_fgate.shape == (B, H, N)
+        log_fgate = log_fgate.float()
+        if has_seq_start:
+            log_fgate = log_fgate.clone()
+            # We absolutely don't want masked value to affect result. If we
+            # don't do this then it could via affecting numerical precision of
+            # cumsum
+            mask_index = (torch.arange(N, device=q.device)[None, None, :] < seq_start[:, None, None])
+            mask_index = torch.broadcast_to(mask_index, log_fgate.size())
+            log_fgate[mask_index] = 0.0
+
+        log_lambda = torch.cumsum(log_fgate, dim=-1, dtype=log_fgate.dtype).float()
+
+        Hk, Hv = k.shape[1], v.shape[1]
+        assert Hk == Hv, "num of heads in k and v should be equal"
+        assert H == Hk, "groupped query attention has not been tested. You can uncomment this if you know what you are doing."
+        assert H % Hk == 0, "number of heads in q must be a multiple of that in k & v"
+        num_groups = H // Hk
+
+        P_SEQ = N - M
+        larger_m = M > N
+        assert (not larger_m), "The key/value tensors must be longer than the query tensor"
+
+        if sm_scale is None:
+            sm_scale = 1. / math.sqrt(D)
+
+        # contiguity
+        q, k, v = maybe_contiguous(q), maybe_contiguous(k), maybe_contiguous(v)
+
+        # to work around https://github.com/openai/triton/issues/2441
+        device = torch.cuda.device_of(q)
+
+        with torch.cuda.device(device):
+
+            config = get_fwd_config(B, H, M, N, D, causal)
+            BLOCK_M, BLOCK_N, num_stages, num_warps = config
+
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+            # consider using 3d grid to avoid div & rem
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            o = torch.empty_like(q)
+            L = torch.empty((B, H, M), device=q.device, dtype=torch.float32)
+            _fwd_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale,
+                L, o,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+                B, H, M, N, P_SEQ, num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=D,
+                IS_CAUSAL=causal, LARGER_M=larger_m, HAS_SEQ_START=has_seq_start,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
+                num_warps=num_warps, num_stages=num_stages,
+            )
+
+        # autograd context maintenance
+        ctx.save_for_backward(q, k, v, o, L, log_lambda, seq_start)
+        ctx.sm_scale = sm_scale
+        ctx.causal = causal
+        ctx.has_seq_start = has_seq_start
+
+        has_extra_return = return_log_normalizer
+        if has_extra_return:
+            outs = (
+                o,
+                L if return_log_normalizer else None,
+            )
+            return outs
+        return o
+
+    @staticmethod
+    def backward(ctx, do, *ignored):
+        q, k, v, o, L, log_lambda, seq_start = ctx.saved_tensors
+        sm_scale = ctx.sm_scale
+        causal = ctx.causal
+        has_seq_start = ctx.has_seq_start
+
+        B, H, M, D = q.shape
+        N = k.shape[2]
+        Hk = k.shape[1]
+        num_groups = H // Hk
+        P_SEQ = N - M
+        larger_m = M > N
+
+        if sm_scale is None:
+            sm_scale = 1. / math.sqrt(D)
+
+        # to work around https://github.com/openai/triton/issues/2441
+        device = torch.cuda.device_of(q)
+        with torch.cuda.device(device):
+            config = get_bwd_config(B, H, M, N, D, causal)
+            BLOCK_M, BLOCK_N, num_stages, num_warps = config
+
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+
+            delta = torch.empty_like(L)
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            _bwd_preprocess[grid](
+                o, do,
+                delta,
+                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                delta.stride(0), delta.stride(1), delta.stride(2),
+                M,
+                BLOCK_M=BLOCK_M, D_HEAD=D,
+                DIVISIBLE_M=divisible_m,
+            )
+
+            # NOTE that dk & dv always have the same number of heads as q, instead of q.
+            BLOCK_M, BLOCK_N, num_stages, num_warps = get_bwd_kv_config(B, H, M, N, D, causal)
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+
+            dk = torch.empty((B, H, N, D), dtype=k.dtype, device=q.device)
+            dv = torch.empty((B, H, N, D), dtype=v.dtype, device=q.device)
+            dlog_lambda = torch.empty((B, H, N), dtype=log_lambda.dtype, device=q.device)
+            grid = (triton.cdiv(N, BLOCK_N), H, B)
+            _bwd_kv_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale, do,
+                dk, dv, dlog_lambda,
+                L, delta,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
+                dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3),
+                dlog_lambda.stride(0), dlog_lambda.stride(1), dlog_lambda.stride(2),
+                B, H, M, N, P_SEQ,
+                num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N, CAUSAL=causal,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n, HAS_SEQ_START=has_seq_start,
+                num_stages=num_stages, num_warps=num_warps,
+            )
+
+            BLOCK_M, BLOCK_N, num_stages, num_warps = get_bwd_q_config(B, H, M, N, D, causal)
+            divisible_m = M % BLOCK_M == 0
+            divisible_n = N % BLOCK_N == 0
+            dq = torch.zeros_like(q)
+            grid = (triton.cdiv(M, BLOCK_M), H, B)
+            _bwd_q_kernel[grid](
+                q, k, v, log_lambda, seq_start, sm_scale, do,
+                dq, dlog_lambda,
+                L, delta,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                log_lambda.stride(0), log_lambda.stride(1), log_lambda.stride(2),
+                do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+                dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
+                dlog_lambda.stride(0), dlog_lambda.stride(1), dlog_lambda.stride(2),
+                B, H, M, N, P_SEQ,
+                num_groups,
+                BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
+                CAUSAL=causal, LARGER_M=larger_m, HAS_SEQ_START=has_seq_start,
+                DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
+                num_stages=num_stages, num_warps = num_warps,
+            )
+            dk = dk.reshape((B, Hk, num_groups, N, D)).sum(2)
+            dv = dv.reshape((B, Hk, num_groups, N, D)).sum(2)
+        dcumsum = torch.cumsum(dlog_lambda, dim=-1, dtype=log_lambda.dtype)
+        dlog_fgate = dlog_lambda + dcumsum[..., -1:] - dcumsum
+        dlog_fgate = dlog_fgate.float()
+        return dq, dk, dv, dlog_fgate, None, None, None, None, None, None, None
+
+
+def forgetting_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_fgate: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+):
+    """
+    A FlashAttention-based implementation of Forgetting Attention. 
+
+    Note:
+    - We recommand bfloat16/float16 for q, k, v and float32 for log_fgate. float32 for 
+      q, k, v is also supported, but the kernel will not use tensor cores if q, k, v are
+      in float32 (which would be slow).
+    - We only support seqlen_q <= seqlen_k
+    - We only support causal attention
+    - Head dimension must be in one of {16, 32, 64, 128}
+
+    Arguments:
+        - q: (batch_size, seqlen_q, num_heads, head_dim) unless head_first=True.
+        - k: (batch_size, seqlen_k, num_heads, head_dim) unless head_first=True.
+        - v: (batch_size, seqlen_k, num_heads, head_dim) unless head_first=True.
+        - log_fgate: (batch_size, seqlen_k, num_heads) unless head_first=True. 
+              This should be the **log** of the forget gates. This is typically the 
+              output of torch.nn.functional.logsigmoid.
+        - head_first: if True, the order the num_heads and seqlen_* axis of the all 
+              FloatTensor inputs and outputs should be (num_heads, seq_len_*) instead of
+              (seq_len_*, num_heads)
+        - seq_start: If not None, should be LongTensor with shape (batch_size,) 
+              and range in [0, seq_len_k). For each batch index batch_id, no attention 
+              will be allocated to tokens before the token index seq_start[batch_id]. 
+              This is useful for left-padded inputs.
+        - sm_scale: The scaling of attention scores before applying softmax. If
+              None, it defaults to (1.0 / math.sqrt(head_dim))
+
+    Returns:
+        out (torch.Tensor): (batch_size, seqlen_q, num_heads, head_dim) unless head_first=True.
+    """
+    if not head_first:
+        q, k, v = [rearrange(item, "b t h d -> b h t d") for item in (q, k, v)]
+        log_fgate = rearrange(log_fgate, "b t h -> b h t")
+    out = ForgettingAttention.apply(q, k, v, log_fgate, seq_start, True, sm_scale, False)
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    return out
+
+
+# --------------------------- Forward ---------------------------
+# NOTE: this function can be overwritten at runtime to use your custom config
+def get_fwd_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0):
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 32, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 4, 4
+    elif torch.cuda.get_device_capability() == (9, 0):
+        # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 2, 8
+    elif torch.cuda.get_device_capability() == (8, 6):
+        if not causal:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+        else: # causal
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 3, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9):
+        # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+
+@triton.jit
+def _fwd_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale,
+    L, O,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr, LARGER_M: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_m = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    log2e: tl.constexpr = 1.4426950408889634
+    loge2: tl.constexpr = 0.6931471805599453
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    O += off_z * stride_oz + off_h * stride_oh
+    L += (off_z * H + off_h) * M # l's shape is (B, H, M)
+
+    offs_m_base = tl.arange(0, BLOCK_M)
+    offs_m = start_m * BLOCK_M + offs_m_base
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m) * stride_log_lambda_n
+    o_ptrs = O + (offs_m[:, None] * stride_om + offs_k[None, :] * stride_ok) # (BLOCK_M, BLOCK_DMODEL)
+    l_ptrs = L + offs_m
+
+    # initialize pointer to m and l, fp32 for accumulators
+    m_i = tl.full([BLOCK_M], value=-float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    # load q
+    if DIVISIBLE_M:
+        q = tl.load(q_ptrs, cache_modifier=".cg")
+        log_lambda_out = tl.load(log_lambda_out_ptrs, cache_modifier=".cg")
+    else:
+        mask_m = offs_m < M
+        q = tl.load(q_ptrs, mask=mask_m[:, None], cache_modifier=".cg")
+        log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m, cache_modifier=".cg")
+
+    #Dot I trick: to place q in registers, it saves shared memory
+    # if BLOCK_DMODEL < 128:
+    #     I = tl.where(offs_k[:, None] == offs_k,
+    #                  tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 1.0, dtype=input_dtype),
+    #                  tl.full((BLOCK_DMODEL, BLOCK_DMODEL), 0.0, dtype=input_dtype))
+    #     q = tl.dot(q, I, input_precision="ieee").to(input_dtype)
+    # else:
+    #     I = tl.where(offs_m_base[:, None] == offs_m_base,
+    #                  tl.full((BLOCK_M, BLOCK_M), 1.0, dtype=input_dtype),
+    #                  tl.full((BLOCK_M, BLOCK_M), 0.0, dtype=input_dtype))
+    #     q = tl.dot(I, q, input_precision="ieee").to(input_dtype)
+
+    # NOTE: Loop-Bound-For-N
+    # The indices in m-dimension that this block may access is in `[start_m * BLOCK_M, (start_m + 1) * BLOCK_M)`.
+    # According to the rule of causal masking, then max index in n-dimension that this block may access
+    # is `P_SEQ + (start_m + 1) * BLOCK_M`.
+    # However, the upper bound of index in n-dimension should never exceed the sequence length of k/v(`P_SEQ + N_CTX`).
+    # `P_SEQ + (start_m + 1) * BLOCK_M` may be larger than `N`.
+    # At this case, there would be illegal memory access when loading k & v tiles
+    # if mask_n is not applied for loading(only when `DIVISIBLE_N`` is true).
+    # See also https://github.com/FlagOpen/FlagAttention/pull/8
+    if IS_CAUSAL:
+        hi = tl.minimum(N, P_SEQ + (start_m + 1) * BLOCK_M)
+        if LARGER_M:
+            hi = tl.maximum(0, hi)
+    else:
+        hi = N
+
+    offs_n_init = offs_n_base
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        lo = tl.minimum(seq_start, hi)
+        lo = (lo // BLOCK_N) * BLOCK_N
+        offs_n_init += lo
+    else:
+        lo = 0
+        seq_start = 0
+
+    # loop over k, v and update accumulators
+    k_ptrs = K + (offs_k[:, None] * stride_kk + offs_n_init[None, :] * stride_kn) # (BLOCK_DMODEL, BLOCK_N)
+    v_ptrs = V + (offs_n_init[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n_init * stride_log_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+    for start_n in range(lo, hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        offs_n = start_n + offs_n_base
+
+        # -- load k, v --
+        if DIVISIBLE_N:
+            k = tl.load(k_ptrs, cache_modifier=".cg")
+            v = tl.load(v_ptrs, cache_modifier=".cg")
+            log_lambda_in = tl.load(log_lambda_in_ptrs, cache_modifier=".cg")
+        else:
+            mask_n = offs_n < N
+            k = tl.load(k_ptrs, mask=mask_n[None, :], cache_modifier=".cg")
+            v = tl.load(v_ptrs, mask=mask_n[:, None], cache_modifier=".cg")
+            log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n, cache_modifier=".cg")
+
+        # -- compute qk ---
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        s = tl.dot(q, k, input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[:, None] - log_lambda_in[None, :]
+        s += decay_bias * log2e
+
+        if not DIVISIBLE_N:
+            s = tl.where(mask_n[None, :], s, float("-inf"))
+        if IS_CAUSAL:
+            causal_mask = (P_SEQ + offs_m[:, None]) >= offs_n[None, :]
+            s = tl.where(causal_mask, s, float("-inf"))
+        if HAS_SEQ_START:
+            s = tl.where(offs_n[None, :] >= seq_start, s, float("-inf"))
+
+
+        # -- compute scaling constant ---
+        m_i_new = tl.maximum(m_i, tl.max(s, 1))
+        alpha = tl.math.exp2((m_i - m_i_new))
+        p = tl.math.exp2(s - m_i_new[:, None])
+
+        # -- compute partial sumexpn before applying dropout
+        p_sum = tl.sum(p, 1)
+
+
+        # -- scale and update acc: acc *= alpha[:, None]--
+        acc *= alpha[:, None]
+        acc += tl.dot(p.to(input_dtype), v, input_precision="ieee")
+
+        # -- update m_i and l_i --
+        l_i = l_i * alpha + p_sum
+        m_i = m_i_new
+        # update pointers
+        k_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vn
+        log_lambda_in_ptrs += BLOCK_N * stride_log_lambda_n
+
+    # write back l & o
+    if IS_CAUSAL and (LARGER_M or HAS_SEQ_START):
+        is_empty_line = (offs_m + P_SEQ) < seq_start
+        acc = tl.where(is_empty_line[:, None], 0.0, acc * (1.0 / l_i[:, None]))
+        l = tl.where(is_empty_line, float("-inf"), m_i * loge2 + tl.log(l_i))
+    else:
+        acc = acc * (1.0 / l_i[:, None])
+        l = m_i * loge2 + tl.log(l_i) # log(normalizer)
+
+
+    if DIVISIBLE_M:
+        tl.store(l_ptrs, l, cache_modifier=".cg")
+        tl.store(o_ptrs, acc.to(input_dtype), cache_modifier=".cg")
+    else:
+        tl.store(l_ptrs, l, mask=mask_m, cache_modifier=".cg")
+        tl.store(o_ptrs, acc.to(input_dtype), mask=mask_m[:, None], cache_modifier=".cg")
+
+
+# --------------------------- Backward ---------------------------
+# NOTE: this function can be overwritten at runtime to use your custom config
+def get_bwd_config(B, H, M, N, D, causal):
+    if torch.cuda.get_device_capability() == (9, 0):
+        if not causal:
+            BLOCK_M = 128 if D <= 64 else 64
+            BLOCK_N = 64
+            num_stages = 2
+            num_warps = 4
+        else:
+            BLOCK_M = 64
+            BLOCK_N = 64
+            num_stages = 3 if D <= 64 else 2
+            num_warps = 4
+    elif torch.cuda.get_device_capability() == (8, 0):
+        if not causal:
+            BLOCK_M = 128 if D <= 64 else 64
+            BLOCK_N = 64
+            num_stages = 2
+            num_warps = 4
+        else:
+            BLOCK_M = 64
+            BLOCK_N = 64
+            num_stages = 3 if D <= 64 else 2
+            num_warps = 4
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if not causal:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 8
+        else:
+            if D <= 64:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+            else:
+                BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 1, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+def get_bwd_kv_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0): # A100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 4, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 128, 4, 8
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9): # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 128, 4, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 128, 2, 8
+    elif torch.cuda.get_device_capability() == (9, 0): # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+def get_bwd_q_config(B, H, M, N, D, causal):
+    assert causal
+    if torch.cuda.get_device_capability() == (8, 0): # A100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 3, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 64, 4, 8
+    elif torch.cuda.get_device_capability() == (8, 6): # tune for RTX-3090, device_capability(8, 6)
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 32, 32, 2, 4
+    elif torch.cuda.get_device_capability() == (8, 9): # L40S
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 4, 4
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 32, 3, 4
+    elif torch.cuda.get_device_capability() == (9, 0): # H100
+        if D <= 64:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 4, 8
+        else:
+            BLOCK_M, BLOCK_N, num_stages, num_warps = 128, 128, 2, 8
+    else:
+        BLOCK_M, BLOCK_N, num_stages, num_warps = 64, 64, 2, 4
+    return (BLOCK_M, BLOCK_N, num_stages, num_warps)
+
+
+@triton.jit
+def _bwd_preprocess(
+    Out, DO,
+    Delta,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dz, stride_dh, stride_dm,
+    M,
+    BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr,
+):
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+    Out += off_z * stride_oz + off_h * stride_oh
+    DO += off_z * stride_doz + off_h * stride_doh
+    Delta += off_z * stride_dz + off_h * stride_dh
+
+    # compute (Out * Dout).sum() for vector interpretation
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_n = tl.arange(0, D_HEAD)
+
+    # load
+    o_ptrs = Out + off_m[:, None] * stride_om + off_n[None, :] * stride_ok
+    do_ptrs = DO + off_m[:, None] * stride_dom + off_n[None, :] * stride_dok
+
+    if DIVISIBLE_M:
+        o = tl.load(o_ptrs).to(tl.float32)
+        do = tl.load(do_ptrs).to(tl.float32)
+    else:
+        mask_m = off_m < M
+        o = tl.load(o_ptrs, mask=mask_m[:, None]).to(tl.float32)
+        do = tl.load(do_ptrs, mask=mask_m[:, None]).to(tl.float32)
+
+    # compute
+    delta = tl.sum(o * do, axis=1)
+
+    # write-back
+    d_ptrs = Delta + off_m * stride_dm
+    if DIVISIBLE_M:
+        tl.store(d_ptrs, delta)
+    else:
+        tl.store(d_ptrs, delta, mask=mask_m)
+
+
+@triton.jit
+def _bwd_kv_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale, DO,
+    DK, DV, DLOG_LAMBDA,
+    L,
+    D,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dkz, stride_dkh, stride_dkn, stride_dkk,
+    stride_dvz, stride_dvh, stride_dvn, stride_dvk,
+    stride_dlog_lambda_z, stride_dlog_lambda_h, stride_dlog_lambda_n,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_n = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+    log2e: tl.constexpr = 1.4426950408889634
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    DO += off_z * stride_doz + off_h * stride_doh
+
+    # offset pointers for batch/head
+    DK += off_z * stride_dkz + off_h * stride_dkh
+    DV += off_z * stride_dvz + off_h * stride_dvh
+    DLOG_LAMBDA += off_z * stride_dlog_lambda_z + off_h * stride_dlog_lambda_h
+
+    # offset pointers for batch/head
+    D += (off_z * H + off_h) * M
+    L += (off_z * H + off_h) * M
+
+    if CAUSAL:
+        lo = tl.maximum(start_n * BLOCK_N - P_SEQ, 0)
+        lo = (lo // BLOCK_M) * BLOCK_M
+    else:
+        lo = 0
+
+    offs_m_init = lo + tl.arange(0, BLOCK_M)
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_base = tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m_init[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m_init) * stride_log_lambda_n # (BLOCK_N, BLOCK_DMODEL)
+    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk) # (BLOCK_N, BLOCK_DMODEL)
+    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n * stride_log_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+    do_ptrs = DO + (offs_m_init[:, None] * stride_dom + offs_k[None, :] * stride_dok) # (BLOCK_M, BLOCK_DMODEL)
+
+    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_k[None, :] * stride_dvk) # (BLOCK_N, BLOCK_DMODEL)
+    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_k[None, :] * stride_dkk) # (BLOCK_N, BLOCK_DMODEL)
+    dlog_lambda_in_ptrs = DLOG_LAMBDA + (offs_n * stride_dlog_lambda_n) # (BLOCK_N, BLOCK_DMODEL)
+
+    # k and v stay in SRAM throughout
+    if DIVISIBLE_N:
+        v = tl.load(v_ptrs)
+        k = tl.load(k_ptrs)
+        log_lambda_in = tl.load(log_lambda_in_ptrs)
+    else:
+        mask_n = offs_n < N
+        v = tl.load(v_ptrs, mask=mask_n[:, None])
+        k = tl.load(k_ptrs, mask=mask_n[:, None])
+        log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n)
+
+    # If the N block doesn't contain seq_start, no need to loop
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        hi = tl.where(start_n * BLOCK_N + BLOCK_N >= seq_start - 1, M, lo)
+    else:
+        hi = M
+
+    # initialize dk amd dv
+    dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dlog_lambda_in = tl.zeros([BLOCK_N], dtype=tl.float32)
+
+    # loop over a col
+    for start_m in range(lo, hi, BLOCK_M):
+        start_m = tl.multiple_of(start_m, BLOCK_M)
+        offs_m = start_m + offs_m_base
+        causal_mask = (P_SEQ + offs_m[None, :]) >= (offs_n[:, None]) # (BLOCK_M, BLOCK_N)
+
+        # load q1, k1, q2, k2, v, do on-chip
+        if DIVISIBLE_M:
+            q = tl.load(q_ptrs)
+            log_lambda_out = tl.load(log_lambda_out_ptrs)
+        else:
+            mask_m = offs_m < M
+            valid_mask = mask_m[None, :] # & mask_n
+            q = tl.load(q_ptrs, mask=mask_m[:, None])
+            log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m)
+        # recompute p = softmax(qk * sm_scale, dim=-1)
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        sT = tl.dot(k, tl.trans(q), input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[None, :] - log_lambda_in[:, None]
+        sT += decay_bias * log2e
+        # NOTE: since softmax in backward is pointwise, the normalizer has been saved in fwd)
+        # So masking on s is not needed.
+        # s = tl.where(valid_mask, s , float("-inf"))
+        # if CAUSAL:
+        #     s = tl.where(causal_mask, s, float("-inf"))
+
+        # -- recompute p ---
+        if DIVISIBLE_M:
+            l = tl.load(L + offs_m)
+        else:
+            l = tl.load(L + offs_m, mask=mask_m)
+        pT = tl.math.exp2(sT - l[None, :] * log2e) # (BLOCK_M, BLOCK_N)
+
+        if not DIVISIBLE_M:
+            pT = tl.where(valid_mask, pT, 0.0)
+        if CAUSAL:
+            pT = tl.where(causal_mask, pT, 0.0)
+
+        # compute dv = dot(p, do)
+        if DIVISIBLE_M:
+            do = tl.load(do_ptrs)
+        else:
+            do = tl.load(do_ptrs, mask=mask_m[:, None]) # (BLOCK_M, BLOCK_DMODEL)
+
+
+        dv += tl.dot(pT.to(input_dtype), do, input_precision="ieee") # (BLOCK_N, BLOCK_DMODEL)  # still correct
+
+        # compute dp = dot(v, do)
+        if DIVISIBLE_M:
+            delta = tl.load(D + offs_m)
+        else:
+            delta = tl.load(D + offs_m, mask=mask_m)
+        # dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        dpT = tl.dot(v, tl.trans(do), input_precision="ieee")
+
+
+        # compute ds = p * (dp - delta[:, None])
+        dsT = pT * (dpT - delta[None, :]) # (BLOCK_M, BLOCK_N)
+
+        if not DIVISIBLE_M:
+            dsT = tl.where(valid_mask, dsT, 0.0)
+        if CAUSAL:
+            dsT = tl.where(causal_mask, dsT, 0.0)
+
+        # compute dk = dot(ds.T, q) masking
+        dk += tl.dot(dsT.to(input_dtype), q, input_precision="ieee")
+        dlog_lambda_in += -tl.sum(dsT, axis=1)
+
+        # increment pointers
+        q_ptrs += BLOCK_M * stride_qm
+        log_lambda_out_ptrs += BLOCK_M * stride_log_lambda_n
+        do_ptrs += BLOCK_M * stride_dom
+
+    dk *= sm_scale
+    if HAS_SEQ_START:
+        # Mask out 
+        seq_mask = (offs_n >= seq_start)
+        dk = tl.where(seq_mask[:, None], dk, 0.0)
+        dv = tl.where(seq_mask[:, None], dv, 0.0)
+        dlog_lambda_in = tl.where(seq_mask, dlog_lambda_in, 0.0)
+    if DIVISIBLE_N:
+        tl.store(dk_ptrs, dk.to(input_dtype)) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dv_ptrs, dv.to(input_dtype)) # (BLOCK_N, BLOCK_DMODEL,)
+        tl.store(dlog_lambda_in_ptrs, dlog_lambda_in.to(tl.float32)) # (BLOCK_N, BLOCK_DMODEL,)
+    else:
+        tl.store(dk_ptrs, dk.to(input_dtype), mask=mask_n[:, None]) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dv_ptrs, dv.to(input_dtype), mask=mask_n[:, None]) # (BLOCK_N, BLOCK_DMODEL)
+        tl.store(dlog_lambda_in_ptrs, dlog_lambda_in.to(tl.float32), mask=mask_n) # (BLOCK_N, BLOCK_DMODEL,)
+
+
+@triton.jit
+def _bwd_q_kernel(
+    Q, K, V, LOG_LAMBDA, SEQ_START, sm_scale, DO,
+    DQ, DLOG_LAMBDA,
+    L,
+    D,
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_log_lambda_z, stride_log_lambda_h, stride_log_lambda_n,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    stride_dqz, stride_dqh, stride_dqm, stride_dqk,
+    stride_dlog_lambda_z, stride_dlog_lambda_h, stride_dlog_lambda_n,
+    Z, H, M, N, P_SEQ,
+    num_groups,
+    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+    CAUSAL: tl.constexpr, LARGER_M: tl.constexpr, HAS_SEQ_START: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
+):
+    input_dtype = Q.dtype.element_ty
+    # -- grid id --
+    start_m = tl.program_id(0)
+    off_h = tl.program_id(1)
+    off_z = tl.program_id(2)
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    log2e: tl.constexpr = 1.4426950408889634
+    qk_scale = sm_scale * log2e
+
+    # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
+    Q += off_z * stride_qz + off_h * stride_qh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
+    LOG_LAMBDA += off_z * stride_log_lambda_z + off_h * stride_log_lambda_h
+    DO += off_z * stride_doz + off_h * stride_doh
+    D += (off_z * H + off_h) * M
+    L += (off_z * H + off_h) * M
+
+    # offset pointers for batch/head
+    DQ += off_z * stride_dqz + off_h * stride_dqh
+    DLOG_LAMBDA += off_z * stride_dlog_lambda_z + off_h * stride_dlog_lambda_h
+
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, BLOCK_DMODEL)
+
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk) # (BLOCK_M, BLOCK_DMODEL)
+    log_lambda_out_ptrs = LOG_LAMBDA + (P_SEQ + offs_m) * stride_log_lambda_n
+
+    dq_ptrs = DQ + (offs_m[:, None] * stride_dqm + offs_k[None, :] * stride_dqk) # (BLOCK_M, BLOCK_DMODEL)
+    dlog_lambda_out_ptrs = DLOG_LAMBDA + (P_SEQ + offs_m) * stride_dlog_lambda_n
+    do_ptrs = DO + (offs_m[:, None] * stride_dom + offs_k[None, :] * stride_dok) # (BLOCK_M, BLOCK_DMODEL)
+
+    # pointer to row-wise quantities in value-like data
+    d_ptrs = D + offs_m
+    l_ptrs = L + offs_m
+
+    # load q: it will stay in SRAM throughout
+    if DIVISIBLE_M:
+        q = tl.load(q_ptrs)
+        do = tl.load(do_ptrs)
+        delta = tl.load(d_ptrs)
+        l = tl.load(l_ptrs)
+        log_lambda_out = tl.load(log_lambda_out_ptrs)
+    else:
+        mask_m = offs_m < M
+        q = tl.load(q_ptrs, mask=mask_m[:, None])
+        do = tl.load(do_ptrs, mask=mask_m[:, None])
+        delta = tl.load(d_ptrs, mask=mask_m)
+        l = tl.load(l_ptrs, mask=mask_m)
+        log_lambda_out = tl.load(log_lambda_out_ptrs, mask=mask_m)
+
+    # initialize dq
+    dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    dlog_lambda_out = tl.zeros([BLOCK_M], dtype=tl.float32)
+
+    # loop over k, v and update accumulator
+    # see note "Loop-Bound-For-N"
+    if CAUSAL:
+        hi = tl.minimum(N, P_SEQ + (start_m + 1) * BLOCK_M)
+        if LARGER_M:
+            hi = tl.maximum(0, hi)
+    else:
+        hi = N
+
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_n_init = offs_n_base
+    if HAS_SEQ_START:
+        SEQ_START += off_z
+        seq_start = tl.load(SEQ_START)
+        lo = tl.minimum(seq_start, hi)
+        lo = (lo // BLOCK_N) * BLOCK_N
+        offs_n_init += lo
+    else:
+        lo = 0
+    k_ptrs = K + (offs_n_init[:, None] * stride_kn + offs_k[None, :] * stride_kk) # (BLOCK_N, BLOCK_DMODEL)
+    v_ptrs = V + (offs_n_init[:, None] * stride_vn + offs_k[None, :] * stride_vk) # (BLOCK_N, BLOCK_DMODEL)
+    log_lambda_in_ptrs = LOG_LAMBDA + (offs_n_init * stride_log_lambda_n)
+
+    # loop over a row
+    for start_n in range(lo, hi, BLOCK_N):
+        offs_n = start_n + offs_n_base
+
+        # load k1, k2, v on chip
+        if DIVISIBLE_N:
+            v = tl.load(v_ptrs)
+            k = tl.load(k_ptrs)
+            log_lambda_in = tl.load(log_lambda_in_ptrs)
+        else:
+            mask_n = offs_n < N
+            v = tl.load(v_ptrs, mask=mask_n[:, None])
+            k = tl.load(k_ptrs, mask=mask_n[:, None])
+            log_lambda_in = tl.load(log_lambda_in_ptrs, mask=mask_n)
+
+
+        # recompute p = softmax(qk * sm_scale, dim=-1)
+        if not DIVISIBLE_N:
+            valid_mask = mask_n[None, :] # & mask_m[:, None]
+        if CAUSAL:
+            causal_mask = (P_SEQ + offs_m[:, None]) >= (offs_n[None, :]) # (BLOCK_M, BLOCK_N)
+        # s = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        s = tl.dot(q, tl.trans(k), input_precision="ieee") * qk_scale
+        decay_bias = log_lambda_out[:, None] - log_lambda_in[None, :]
+        s += decay_bias * log2e
+
+        # NOTE: since softmax in backward is pointwise, the normalizer has been saved in fwd)
+        # So masking on s is not needed.
+        # if CAUSAL:
+        #     s = tl.where(causal_mask & valid_mask, s, float("-inf"))
+        # else:
+        #     s = tl.where(valid_mask, s, float("-inf"))
+        p = tl.math.exp2(s - l[:, None] * log2e) # (BLOCK_M, BLOCK_N)
+
+        # compute dp = dot(v, do)
+        # dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        dp = tl.dot(do.to(input_dtype), tl.trans(v), input_precision="ieee")
+
+
+        # no need to mask dp
+        # if CAUSAL:
+        #     dp = tl.where(causal_mask & valid_mask, dp, 0.0)
+        # else:
+        #     dp = tl.where(valid_mask, dp, 0.0)
+
+        # compute ds = p * (dp - delta[:, None])
+        # move scale out to dq at last
+        ds = p * (dp - delta[:, None]) # (BLOCK_M, BLOCK_N)
+
+        # mask ds to ensure no small values
+        if not DIVISIBLE_N:
+            ds = tl.where(valid_mask, ds, 0.0)
+        if CAUSAL:
+            ds = tl.where(causal_mask, ds, 0.0)
+        if HAS_SEQ_START:
+            ds = tl.where(offs_n[None, :] >= seq_start, ds, 0.0)
+
+        dq += tl.dot(ds.to(input_dtype), k, input_precision="ieee")
+        dlog_lambda_out += tl.sum(ds, axis=1)
+
+        # increment pointers
+        k_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vn
+        log_lambda_in_ptrs += BLOCK_N * stride_log_lambda_n
+
+    dq *= sm_scale
+    if DIVISIBLE_M:
+        tmp = tl.load(dlog_lambda_out_ptrs)
+    else:
+        tmp = tl.load(dlog_lambda_out_ptrs, mask=mask_m)
+    dlog_lambda_out += tmp
+    if DIVISIBLE_M:
+        tl.store(dq_ptrs, dq.to(input_dtype))
+        tl.store(dlog_lambda_out_ptrs, dlog_lambda_out)
+    else:
+        tl.store(dq_ptrs, dq.to(input_dtype), mask=mask_m[:, None])
+        tl.store(dlog_lambda_out_ptrs, dlog_lambda_out, mask=mask_m)
+
+
+
+@pytest.mark.parametrize("Z, H, M, N, HEAD_DIM", [(4, 2, 1020, 2098, 64), (4, 2, 1024, 2048, 64)])
+@pytest.mark.parametrize("causal", [True])
+def test_op(Z, H, M, N, HEAD_DIM, causal, dtype=torch.bfloat16):
+    torch.manual_seed(24)
+    q = (torch.empty((Z, H, M, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    k = (torch.empty((Z, H, N, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    v = (torch.empty((Z, H, N, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
+    fgate_logit = torch.empty((Z, H, N), dtype=torch.float32, device="cuda").uniform_(5, 10)
+    log_fgate = torch.nn.functional.logsigmoid(fgate_logit).requires_grad_()
+    seq_start = torch.randint(low=0, high=N, size=(Z,), dtype=torch.long, device="cuda")
+    # seq_start = torch.randint(low=0, high=10, size=(Z,), dtype=torch.long, device="cuda")
+    # seq_start = torch.full(fill_value=0, size=(Z,), dtype=torch.long, device="cuda")
+    sm_scale = 0.5
+    dout = torch.randn_like(q)
+    # reference implementation
+    P_SEQ = N - M
+    mask = torch.tril(torch.ones((M, N), device="cuda"), diagonal=P_SEQ)
+    p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
+    p = p.float()
+
+    log_lambda = torch.cumsum(log_fgate, dim=-1)
+    decay_bias = log_lambda[..., -M:, None] - log_lambda[..., None, :]
+    p = p + decay_bias
+    if causal:
+        p[:, :, mask == 0] = float("-inf")
+
+    attention_mask = torch.arange(N, device="cuda") < seq_start[:, None, None, None]
+    p = torch.where(attention_mask, float("-inf"), p)
+    p = torch.softmax(p.float(), dim=-1).to(dtype)
+    p = p.clone()
+    p[torch.isnan(p)] = 0.0
+    # p = torch.exp(p)
+    ref_out = torch.matmul(p, v)
+    ref_out.backward(dout)
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dlog_fgate, log_fgate.grad = log_fgate.grad.clone(), None
+    # triton implementation
+    tri_out = forgetting_attention(q, k, v, log_fgate, head_first=True, seq_start=seq_start, sm_scale=sm_scale)
+    tri_out = tri_out.to(dtype)
+
+    tri_out.backward(dout)
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dlog_fgate, log_fgate.grad = log_fgate.grad.clone(), None
+    # compare
+    # assert torch.allclose(tri_log_normalizer[~torch.isnan(tri_log_normalizer)], ref_log_normalizer[~torch.isnan(ref_log_normalizer)], atol=1e-2, rtol=0)
+    assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0), (ref_out - tri_out).abs().max()
+    rtol = 0
+    # Relative tolerance workaround for known hardware limitation of MI200 GPU.
+    # For details see https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
+    # if torch.version.hip is not None and triton.runtime.driver.active.get_current_target().arch == "gfx90a":
+        # rtol = 1e-2
+    assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=rtol), (ref_dv - tri_dv).abs().max()
+    assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=rtol), (ref_dk - tri_dk).abs().max()
+    assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=rtol), (ref_dq - tri_dq).abs().max()
+    assert torch.allclose(ref_dlog_fgate, tri_dlog_fgate, atol=1e-2, rtol=rtol), (ref_dlog_fgate - tri_dlog_fgate).abs().max()
+
+try:
+    from flash_attn.flash_attn_interface import \
+        flash_attn_qkvpacked_func as flash_attn_func
+    HAS_FLASH = True
+except BaseException:
+    HAS_FLASH = False
+
+TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
+BATCH, N_HEADS, HEAD_DIM = 4, 32, 128
+# vary seq length for fixed head and batch=4
+configs = []
+for mode in ["fwd", "bwd"]:
+# for mode in ["bwd"]:
+    # for causal in [True, False]:
+    for causal in [True]:
+        if mode == "bwd" and not causal:
+            continue
+        configs.append(
+            triton.testing.Benchmark(
+                x_names=["N_CTX"],
+                # x_vals=[2**i for i in range(10, 15)],
+                x_vals=[2**i for i in range(14, 15)],
+                line_arg="provider",
+                # line_vals=["triton-fp16", "flag"] + (["flash"] if HAS_FLASH else []),
+                # line_names=["Triton [FP16]", "Flag"] + (["Flash-2"] if HAS_FLASH else []),
+                line_vals=["flag"] + (["flash"] if HAS_FLASH else []),
+                line_names=["Flag"] + (["Flash-2"] if HAS_FLASH else []),
+                styles=[("red", "-"), ("blue", "-"), ("green", "-")],
+                ylabel="ms",
+                plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}",
+                args={
+                    "H": N_HEADS,
+                    "BATCH": BATCH,
+                    "HEAD_DIM": HEAD_DIM,
+                    "mode": mode,
+                    "causal": causal,
+                },
+            ))
+
+
+@triton.testing.perf_report(configs)
+def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, causal, mode, provider, device="cuda"):
+    assert mode in ["fwd", "bwd"]
+    warmup = 25
+    rep = 100
+    dtype = torch.bfloat16
+    if "flag" in provider:
+        q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fgate_logit = torch.empty((BATCH, H, N_CTX), dtype=torch.float32, device="cuda").uniform_(5, 10)
+        log_fgate = torch.nn.functional.logsigmoid(fgate_logit).requires_grad_()
+        # if mode == "fwd" and "fp8" in provider:
+        #     q = q.to(torch.float8_e5m2)
+        #     k = k.to(torch.float8_e5m2)
+        #     v = v.permute(0, 1, 3, 2).contiguous()
+        #     v = v.permute(0, 1, 3, 2)
+        #     v = v.to(torch.float8_e5m2)
+        sm_scale = 1.3
+        fn = lambda: forgetting_attention(q, k, v, log_fgate, head_first=True, sm_scale=sm_scale)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    if provider == "flash":
+        qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fn = lambda: flash_attn_func(qkv, causal=causal)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM
+    total_flops = 2 * flops_per_matmul
+    if causal:
+        total_flops *= 0.5
+    if mode == "bwd":
+        total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
+    return total_flops / ms * 1e-9
+
+
+if __name__ == "__main__":
+    # only works on post-Ampere GPUs right now
+    bench_flash_attention.run(save_path=".", print_data=True)

ops/forgetting_attention_std.py

@@ -0,0 +1,72 @@
+"""
+Forgetting Attention - 标准 Softmax 版本
+在 forgetting_attention.py 最后添加这个函数
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def forgetting_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_fgate: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """标准 Softmax 版本的 Forgetting Attention"""
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+        log_fgate = rearrange(log_fgate, "b t h -> b h t")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # 计算 QK 分数
+    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # 处理 seq_start
+    log_fgate_masked = log_fgate.float()
+    if seq_start is not None:
+        log_fgate_masked = log_fgate_masked.clone()
+        mask_idx = torch.arange(T_k, device=q.device)[None, None, :] < seq_start[:, None, None]
+        log_fgate_masked[mask_idx] = 0.0
+    
+    # 计算累积衰减
+    log_lambda = torch.cumsum(log_fgate_masked, dim=-1)
+    decay_bias = log_lambda[:, :, :T_q, None] - log_lambda[:, :, None, :]
+    scores = scores + decay_bias
+    
+    # Causal mask
+    P_SEQ = T_k - T_q
+    causal_mask = torch.triu(torch.ones((T_q, T_k), dtype=torch.bool, device=q.device), diagonal=P_SEQ + 1)
+    scores = scores.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # seq_start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        scores = scores.masked_fill(seq_mask, float('-inf'))
+    
+    # Softmax
+    attn = F.softmax(scores, dim=-1)
+    attn = torch.nan_to_num(attn, 0.0)
+    
+    # 计算输出
+    out = torch.matmul(attn.to(v.dtype), v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out

ops/framework_mock.py

@@ -0,0 +1,25 @@
+"""
+Mock framework module for ndr geometric attention
+只保留必要的部分
+"""
+import torch
+from typing import Optional, Any
+
+class visualize:
+    """Mock visualize class"""
+    @staticmethod
+    def attention(*args, **kwargs):
+        """Dummy attention visualization"""
+        pass
+    
+    @staticmethod
+    def plot(*args, **kwargs):
+        """Dummy plot"""
+        pass
+
+# Mock其他可能需要的功能
+def get_logger(name: str):
+    """Mock logger"""
+    import logging
+    return logging.getLogger(name)
+

ops/geometric_attention/__init__.py

@@ -0,0 +1 @@
+from .cuda_interface import geometric_attention_activation

ops/geometric_attention/cuda_interface.cu

@@ -0,0 +1,177 @@
+#include <torch/extension.h>
+
+__global__ void k_cuda_log_sigmoid_forward(int N, float * t, float *out_sigm, float *out_one_minus_sigm){
+    int i = threadIdx.x + blockIdx.x * blockDim.x;
+    if (i<N){
+        float x = t[i];
+        float c = - log(exp(-abs(x)) + 1);
+        out_sigm[i] = min(x, 0.0f) + c;
+        out_one_minus_sigm[i] = -max(x, 0.0f) + c;
+    }
+}
+
+__global__ void k_cuda_log_sigmoid_backward(int N, float *t, float *grad_sigm, float *grad_one_minus_sigm, float *grad_out){
+    int i = threadIdx.x + blockIdx.x * blockDim.x;
+    if (i<N){
+        float x = t[i];
+        float ne = exp(-abs(x));
+        float coeff = 1.0 / (ne + 1.0) * ne;
+
+        float r_one_minus = (x > 0) ? (coeff - 1) : (-coeff);
+        float r =  (x < 0) ? (coeff - 1) : (-coeff);
+        grad_out[i] = - grad_sigm[i] * r + grad_one_minus_sigm[i] * r_one_minus;
+    }
+}
+
+std::vector<torch::Tensor> cuda_log_sigmoid_forward(torch::Tensor input){
+    auto o1 = torch::empty_like(input);
+    auto o2 = torch::empty_like(input);
+    auto inf = input.flatten();
+
+    const int N = inf.size(0);
+
+    const int threads = 256;
+    const int blocks = (N + threads - 1) / threads;
+
+    k_cuda_log_sigmoid_forward<<<blocks, threads>>>(N,  
+        input.data<float>(),
+        o1.data<float>(),
+        o2.data<float>());
+    
+    return {o1, o2};
+}
+
+std::vector<torch::Tensor> cuda_log_sigmoid_backward(torch::Tensor input, torch::Tensor grad_sigm, torch::Tensor grad_one_minus_sigm){
+    auto output = torch::empty_like(input);
+    auto N = input.flatten().size(0);
+
+    const int threads = 256;
+    const int blocks = (N + threads - 1) / threads;
+
+    k_cuda_log_sigmoid_backward<<<blocks, threads>>>(N,  
+        input.data<float>(),
+        grad_sigm.data<float>(),
+        grad_one_minus_sigm.data<float>(),
+        output.data<float>());
+    
+    return {output};
+}
+
+
+typedef torch::PackedTensorAccessor32<float, 3, torch::RestrictPtrTraits> float_accessor;
+
+__global__ void k_cuda_window_sum_forward(float_accessor csum, float_accessor out, int offset){
+    const int in_p = threadIdx.z + blockIdx.z * blockDim.z;
+    const int out_p_mem = threadIdx.y + blockIdx.y * blockDim.y;
+    const int batch = threadIdx.x + blockIdx.x * blockDim.x;
+
+    const int out_p = out_p_mem + offset;
+
+    if (batch < out.size(0) & out_p_mem < out.size(1) & in_p < out.size(2)){
+        float res;
+        if (in_p == out_p){
+            res = 0;
+        } else {
+            const int offset = abs(out_p - in_p);
+            int p_i = out_p + offset - int(in_p > out_p);
+            const int n_i = out_p - offset;
+
+            p_i = min(p_i, out.size(2) - 1);
+
+            float d_n = (n_i >= 0) ? (csum[batch][out_p_mem][n_i]) : 0.0;
+            res = (csum[batch][out_p_mem][p_i]) - d_n;
+        }
+
+        out[batch][out_p_mem][in_p] = res;
+    }
+
+}
+
+__global__ void k_cuda_window_sum_backward(float_accessor grad_in, float_accessor grad_out, int offset){
+    const int in_p = threadIdx.z + blockIdx.z * blockDim.z;
+    const int out_p_mem = threadIdx.y + blockIdx.y * blockDim.y;
+    const int batch = threadIdx.x + blockIdx.x * blockDim.x;
+
+    const int out_p = out_p_mem + offset;
+
+    if (batch < grad_out.size(0) & out_p_mem < grad_out.size(1) & in_p < grad_out.size(2)){
+        const int other = 2 * out_p - in_p;
+
+        float res;
+        if (in_p == grad_out.size(2) - 1){
+            res = 0;
+            for (int i = 0; i < other + int(in_p != out_p); ++i){
+                res += grad_in[batch][out_p_mem][i];
+            }
+        } else if (in_p == out_p){
+            res = grad_in[batch][out_p_mem][min(in_p + 1, grad_out.size(2) - 1)];
+        } else if (in_p < out_p){
+            res = -grad_in[batch][out_p_mem][in_p];
+            if (other < grad_in.size(2))
+                res -= grad_in[batch][out_p_mem][other];
+        } else {
+            res = grad_in[batch][out_p_mem][in_p + 1];
+            if (other >= 0)
+                res += grad_in[batch][out_p_mem][other];
+        }
+
+        grad_out[batch][out_p_mem][in_p] = res;
+    }
+}
+
+dim3 get_grid_size(torch::Tensor target, dim3 block_dim){
+    return dim3(
+        (target.size(0) + block_dim.x - 1) / block_dim.x,
+        (target.size(1) + block_dim.y - 1) / block_dim.y,
+        (target.size(2) + block_dim.z - 1) / block_dim.z
+    );
+}
+
+torch::Tensor cuda_window_sum_forward(torch::Tensor input, int offset){
+    auto out = torch::empty_like(input);
+
+    dim3 block_size(2, 2, 32);
+    k_cuda_window_sum_forward<<<get_grid_size(input, block_size), block_size>>>(
+        input.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
+        out.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
+        offset
+    );
+    
+    return out;
+}
+
+torch::Tensor cuda_window_sum_backward(torch::Tensor grad_in, int offset){
+    auto out = torch::empty_like(grad_in);
+
+    dim3 block_size(2, 2, 32);
+    k_cuda_window_sum_backward<<<get_grid_size(grad_in, block_size), block_size>>>(
+        grad_in.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
+        out.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
+        offset
+    );
+    
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def(
+        "cuda_log_sigmoid_forward",
+        &cuda_log_sigmoid_forward,
+        "Log sigmoid, forward pass"
+    );
+    m.def(
+        "cuda_log_sigmoid_backward",
+        &cuda_log_sigmoid_backward,
+        "Log sigmoid, backward pass"
+    );
+    m.def(
+        "cuda_window_sum_forward",
+        &cuda_window_sum_forward,
+        "Window sum, forward pass"
+    );
+    m.def(
+        "cuda_window_sum_backward",
+        &cuda_window_sum_backward,
+        "Window sum, backward pass"
+    );
+}

ops/geometric_attention/cuda_interface.py

@@ -0,0 +1,93 @@
+import os
+import torch
+import multiprocessing
+from typing import Tuple, Optional
+import torch.nn.functional as F
+import filelock  # 用filelock替代framework.utils.LockFile
+
+# Just in time import
+# https://pytorch.org/tutorials/advanced/cpp_extension
+
+dirname = os.path.dirname(__file__)
+filename = os.path.join(dirname, 'cuda_interface.cu')
+outdir = "./cache/geometric_attention"
+os.makedirs(outdir, exist_ok=True)
+
+cuda_log_sigmoid_backward = None
+cuda_log_sigmoid_forward = None
+cuda_window_sum_forward = None
+cuda_window_sum_backward = None
+
+def load_extension():
+    global cuda_log_sigmoid_forward, cuda_log_sigmoid_backward
+    global cuda_window_sum_forward, cuda_window_sum_backward
+    if cuda_log_sigmoid_forward is not None:
+        return
+
+    # 使用filelock替代framework.utils.LockFile
+    lock = filelock.FileLock(outdir + "/lock.lock")
+    with lock:
+        from torch.utils.cpp_extension import load
+
+        os.environ["MAX_JOBS"] = str(multiprocessing.cpu_count())
+        ext = load(
+            extra_cuda_cflags=['--ftemplate-depth=1024'],
+            name="geometric_attention_cuda_interface",
+            sources=[filename], verbose=True)
+
+        cuda_log_sigmoid_forward = ext.cuda_log_sigmoid_forward
+        cuda_log_sigmoid_backward = ext.cuda_log_sigmoid_backward
+        cuda_window_sum_forward = ext.cuda_window_sum_forward
+        cuda_window_sum_backward = ext.cuda_window_sum_backward
+
+
+class LogSigmoidFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = x.detach().contiguous()
+        ctx.save_for_backward(x)
+        a, b = cuda_log_sigmoid_forward(x)
+        return a, b
+
+    @staticmethod
+    def backward(ctx, grad_in_sigm: torch.Tensor, grad_in_one_minus: torch.tensor) -> torch.Tensor:
+        xf, = ctx.saved_tensors
+        ga = grad_in_sigm.contiguous()
+        gb = grad_in_one_minus.contiguous()
+        return cuda_log_sigmoid_backward(xf, ga, gb)[0]
+
+
+class WindowSumFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, csum: torch.Tensor, offset: int) -> torch.Tensor:
+        ctx.saved_offset = offset
+        c2 = csum.detach().contiguous().flatten(end_dim=-3)
+        res = cuda_window_sum_forward(c2, offset)
+        return res.view_as(csum)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
+        offset = ctx.saved_offset
+        go = grad_output.contiguous().flatten(end_dim=-3)
+        res = cuda_window_sum_backward(go, offset)
+        return res.view_as(grad_output), None
+
+
+def window_sum(x: torch.Tensor, offset: int) -> torch.Tensor:
+    load_extension()
+    return WindowSumFunction.apply(x, offset)
+
+
+def log_sigmoid(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    load_extension()
+    return LogSigmoidFunction.apply(x)
+
+
+def geometric_attention_activation(logits: torch.Tensor, mask: Optional[torch.Tensor] = None, pos_offset: int = 0,
+                                   normalize: bool = True) -> torch.Tensor:
+        p, one_minus_p = log_sigmoid(logits)
+        not_previos = window_sum(one_minus_p.cumsum(-1), pos_offset)
+
+        probs = (not_previos + p).exp()
+
+        return F.normalize(probs, 1, -1) if normalize else probs

ops/geometric_attention/cuda_interface.py.bak

@@ -0,0 +1,94 @@
+import os
+import torch
+import multiprocessing
+from framework.utils import LockFile
+from typing import Tuple, Optional
+import torch.nn.functional as F
+
+# Just in time import
+# https://pytorch.org/tutorials/advanced/cpp_extens
+
+dirname = os.path.dirname(__file__)
+filename = os.path.join(dirname, 'cuda_interface.cu')
+outdir = "./cache/geometric_attention"
+os.makedirs(outdir, exist_ok=True)
+
+cuda_log_sigmoid_backward = None
+cuda_log_sigmoid_forward = None
+cuda_window_sum_forward = None
+cuda_window_sum_backward = None
+
+def load_extension():
+    global cuda_log_sigmoid_forward, cuda_log_sigmoid_backward
+    global cuda_window_sum_forward, cuda_window_sum_backward
+    if cuda_log_sigmoid_forward is not None:
+        return
+
+    with LockFile(outdir + "/lock"):
+        from torch.utils.cpp_extension import load
+
+        os.environ["MAX_JOBS"] = str(multiprocessing.cpu_count())
+        ext = load(
+            extra_cuda_cflags=['--ftemplate-depth=1024'],
+            name="geometric_attention_cuda_interface",
+            sources=[filename], verbose=True)
+            #, build_directory=outdir)
+
+        cuda_log_sigmoid_forward = ext.cuda_log_sigmoid_forward
+        cuda_log_sigmoid_backward = ext.cuda_log_sigmoid_backward
+        cuda_window_sum_forward = ext.cuda_window_sum_forward
+        cuda_window_sum_backward = ext.cuda_window_sum_backward
+
+
+class LogSigmoidFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = x.detach().contiguous()
+        ctx.save_for_backward(x)
+        a, b = cuda_log_sigmoid_forward(x)
+        return a, b
+        # return res_a.view_as(x), res_b.view_as(x)
+
+    @staticmethod
+    def backward(ctx, grad_in_sigm: torch.Tensor, grad_in_one_minus: torch.tensor) -> torch.Tensor:
+        xf, = ctx.saved_tensors
+        ga = grad_in_sigm.contiguous()
+        gb = grad_in_one_minus.contiguous()
+        return cuda_log_sigmoid_backward(xf, ga, gb)[0]
+
+
+class WindowSumFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, csum: torch.Tensor, offset: int) -> torch.Tensor:
+        ctx.saved_offset = offset
+        c2 = csum.detach().contiguous().flatten(end_dim=-3)
+        res = cuda_window_sum_forward(c2, offset)
+        return res.view_as(csum)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        offset = ctx.saved_offset
+        go = grad_output.contiguous().flatten(end_dim=-3)
+        res = cuda_window_sum_backward(go, offset)
+        return res.view_as(grad_output), None
+
+
+def window_sum(x: torch.Tensor, offset: int) -> torch.Tensor:
+    load_extension()
+    return WindowSumFunction.apply(x, offset)
+
+
+def log_sigmoid(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    load_extension()
+    return LogSigmoidFunction.apply(x)
+
+
+def geometric_attention_activation(logits: torch.Tensor, mask: Optional[torch.Tensor] = None, pos_offset: int = 0,
+                                   normalize: bool = True) -> torch.Tensor:
+        p, one_minus_p = log_sigmoid(logits)
+        not_previos = window_sum(one_minus_p.cumsum(-1), pos_offset)
+
+        probs = (not_previos + p).exp()
+
+        # return probs
+        return F.normalize(probs, 1, -1) if normalize else probs

ops/geometric_attention_final.py

@@ -0,0 +1,109 @@
+"""
+Geometric Attention - CUDA加速版本 (支持FP16)
+"""
+
+import math
+import torch
+from einops import rearrange
+from typing import Optional
+
+# 尝试导入CUDA版本
+try:
+    from forgetting_transformer.ops.geometric_attention.cuda_interface import (
+        load_extension, 
+        geometric_attention_activation,
+    )
+    load_extension()
+    HAS_CUDA = True
+    print("✅ Using CUDA geometric attention (with FP16 support)")
+except Exception as e:
+    HAS_CUDA = False
+    print(f"⚠️  CUDA not available: {e}")
+
+
+def geometric_attention_cuda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+) -> torch.Tensor:
+    if not HAS_CUDA:
+        raise RuntimeError("CUDA not available")
+    
+    # ⭐ 保存原始dtype
+    original_dtype = q.dtype
+    needs_cast = original_dtype == torch.float16
+    
+    # ⭐ 如果是FP16，转成FP32
+    if needs_cast:
+        q = q.float()
+        k = k.float()
+        v = v.float()
+    
+    # Rearrange
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Attention scores
+    logits = torch.matmul(q, k.transpose(-2, -1)) * sm_scale
+    
+    # CUDA kernel (FP32)
+    attn_weights = geometric_attention_activation(
+        logits, mask=None, pos_offset=0, normalize=normalize
+    )
+    
+    # Apply to values
+    output = torch.matmul(attn_weights, v)
+    
+    # Rearrange back
+    if not head_first:
+        output = rearrange(output, "b h t d -> b t h d")
+    
+    # ⭐ 转回原始dtype
+    if needs_cast:
+        output = output.to(original_dtype)
+    
+    return output
+
+
+def geometric_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """自动选择CUDA或Python"""
+    
+    if HAS_CUDA and q.is_cuda:
+        try:
+            return geometric_attention_cuda(
+                q, k, v, head_first=head_first,
+                seq_start=seq_start, sm_scale=sm_scale,
+                normalize=normalize
+            )
+        except Exception as e:
+            # 不打印太多警告，会刷屏
+            pass
+    
+    # Fallback
+    from forgetting_transformer.ops.geometric_attention_std import geometric_attention_std
+    return geometric_attention_std(
+        q, k, v, head_first=head_first,
+        seq_start=seq_start, sm_scale=sm_scale,
+        normalize=normalize
+    )

ops/geometric_attention_std.py

@@ -0,0 +1,179 @@
+"""
+Geometric Attention - 标准 Softmax 版本
+基于论文 "The Neural Data Router" (Csordás et al., 2022)
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def geometric_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    标准 Softmax 版本的 Geometric Attention
+    
+    Args:
+        q: Query tensor [B, T, H, D] or [B, H, T, D] if head_first
+        k: Key tensor [B, T, H, D] or [B, H, T, D] if head_first
+        v: Value tensor [B, T, H, D] or [B, H, T, D] if head_first
+        head_first: 是否head维度在前
+        seq_start: 序列起始位置 [B]
+        sm_scale: scaling factor，默认 1/sqrt(D)
+        normalize: 是否归一化attention weights
+    
+    Returns:
+        output: [B, T, H, D] or [B, H, T, D] if head_first
+    """
+    
+    # Rearrange to head_first format
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Step 1: 计算 content-based logits
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    # logits: [B, H, T_q, T_k]
+    
+    # Step 2: Mask diagonal (不允许attend到自己)
+    if T_q == T_k:
+        diag_mask = torch.eye(T_q, dtype=torch.bool, device=q.device)
+        logits = logits.masked_fill(diag_mask[None, None, :, :], float('-inf'))
+    
+    # Step 3: 处理 seq_start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        logits = logits.masked_fill(seq_mask, float('-inf'))
+    
+    # Step 4: Causal mask (如果需要)
+    # 注意：geometric attention论文中没有causal，如果你的任务需要可以取消注释
+    # P_SEQ = T_k - T_q
+    # causal_mask = torch.triu(torch.ones((T_q, T_k), dtype=torch.bool, device=q.device), diagonal=P_SEQ + 1)
+    # logits = logits.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # Step 5: Geometric weighting (核心算法)
+    attn_weights = geometric_weighting(logits, normalize=normalize)
+    
+    # Step 6: 应用attention到values
+    out = torch.matmul(attn_weights.to(v.dtype), v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out
+
+
+def geometric_weighting(
+    logits: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    计算geometric attention weights
+    
+    实现论文中的 Equation 7:
+    A[i,j] = P[i,j] * ∏(1 - P[i,k]) for k closer to i than j
+    
+    Args:
+        logits: [B, H, T_q, T_k] attention logits
+        normalize: 是否归一化
+    
+    Returns:
+        weights: [B, H, T_q, T_k] attention weights
+    """
+    B, H, T_q, T_k = logits.shape
+    
+    # Step 1: Sigmoid to get matching probabilities
+    P = torch.sigmoid(logits)  # [B, H, T_q, T_k]
+    
+    # Step 2: 使用 log-space 计算（数值稳定）
+    log_P = torch.log(P + 1e-10)
+    log_one_minus_P = torch.log(1.0 - P + 1e-10)
+    
+    # Step 3: 简化版本 - 使用cumsum实现几何分布
+    # 这是一个高效的近似，避免了显式的循环
+    
+    # 对于每个位置i，计算其左侧所有位置的log(1-P)累积和
+    log_decay_left = log_one_minus_P.cumsum(dim=-1)
+    
+    # 计算weights（简化版）
+    # 完整版本需要根据距离动态选择区间，这里用一个高效近似
+    weights = torch.exp(log_P + log_decay_left.roll(1, dims=-1))
+    
+    # 第一个位置特殊处理（没有左侧元素）
+    # 避免inplace操作
+    weights_first = P[:, :, :, :1]  # 获取第一列
+    weights = torch.cat([weights_first, weights[:, :, :, 1:]], dim=-1)
+    
+    # Step 4: 归一化（可选）
+    if normalize:
+        weights = F.normalize(weights, p=1, dim=-1)
+    
+    # 处理NaN（如果所有位置都是-inf）
+    weights = torch.nan_to_num(weights, 0.0)
+    
+    return weights
+
+
+def geometric_weighting_full(
+    logits: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    完整版geometric weighting（更慢但更准确）
+    
+    仅在需要最高精度时使用，训练时建议用上面的简化版
+    """
+    B, H, T_q, T_k = logits.shape
+    device = logits.device
+    
+    P = torch.sigmoid(logits)
+    log_P = torch.log(P + 1e-10)
+    log_one_minus_P = torch.log(1.0 - P + 1e-10)
+    
+    # 初始化weights
+    weights = torch.zeros_like(P)
+    
+    # 对每个(i,j)计算geometric weight
+    for i in range(T_q):
+        for j in range(T_k):
+            # 找出比j更接近i的所有位���k
+            if i < j:
+                # 向右看：closer positions are [i+1, ..., j-1]
+                closer_positions = range(i + 1, j)
+            elif i > j:
+                # 向左看：closer positions are [j+1, ..., i-1]
+                closer_positions = range(j + 1, i)
+            else:
+                # i == j (对角线)，已经在外面mask掉了
+                continue
+            
+            # 计算 ∏(1 - P[i,k]) in log-space
+            log_prod = sum(log_one_minus_P[:, :, i, k] for k in closer_positions) if closer_positions else 0.0
+            
+            # weights[i,j] = P[i,j] * ∏(1 - P[i,k])
+            weights[:, :, i, j] = torch.exp(log_P[:, :, i, j] + log_prod)
+    
+    if normalize:
+        weights = F.normalize(weights, p=1, dim=-1)
+    
+    weights = torch.nan_to_num(weights, 0.0)
+    
+    return weights 

ops/layer_with_visualization.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn
+from typing import Dict, Any
+
+
+class LayerWithVisualization(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.visualization_enabled = False
+
+    def prepare(self):
+        # Should be called before the training step
+        pass
+
+    def plot(self, options: Dict[str, Any]) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+
+class LayerVisualizer:
+    def __init__(self, module: torch.nn.Module, options: Dict[str, Any] = {}):
+        self.modules = []
+        self.options = options
+        self.curr_options = None
+        for n, m in module.named_modules():
+            if isinstance(m, LayerWithVisualization):
+                self.modules.append((n, m))
+
+    def plot(self) -> Dict[str, Any]:
+        res = {}
+        for n, m in self.modules:
+            res.update({f"{n}/{k}": v for k, v in m.plot(self.curr_options).items()})
+            m.visualization_enabled = False
+            
+        self.curr_options = None
+        return res
+
+    def prepare(self, options: Dict[str, Any] = {}):
+        self.curr_options = self.options.copy()
+        self.curr_options.update(options)
+
+        for _, m in self.modules:
+            m.prepare()
+            m.visualization_enabled = True

ops/multi_head_attention.py

@@ -0,0 +1,149 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Callable, List, Union, Tuple, Dict, Any
+from dataclasses import dataclass
+from forgetting_transformer.ops.layer_with_visualization import LayerWithVisualization
+import forgetting_transformer.ops.framework_mock as framework
+
+
+@dataclass
+class AttentionMask:
+    src_length_mask: Optional[torch.Tensor]
+    position_mask: Optional[torch.Tensor]
+
+
+class MultiHeadAttentionBase(LayerWithVisualization):
+    def __init__(self, state_size: int, n_heads: int, dropout: float=0.1, projection_size: Optional[int] = None):
+        assert state_size % n_heads == 0
+        super().__init__()
+        self.attention_to_visualize = []
+
+        self.state_size = state_size
+        self.projection_size = projection_size or (state_size // n_heads)
+        self.n_heads = n_heads
+        self.scale = 1.0 / math.sqrt(self.projection_size)
+
+        self.dropout = torch.nn.Dropout(dropout)
+
+    @staticmethod
+    def apply_logit_masks(logits: torch.Tensor, mask: Optional[AttentionMask], val: float = float("-inf")) -> torch.Tensor:
+        if mask.position_mask is not None:
+            # [..., N_out, N_in], broadcast works
+            logits = logits.masked_fill(mask.position_mask, val)
+
+        if mask.src_length_mask is not None:
+            # [B, ...., N_in], needs manual shaping
+            b, i = mask.src_length_mask.shape
+            pad_dims = logits.ndim - 2
+            logits = logits.masked_fill(mask.src_length_mask.view([b] + [1] * pad_dims + [i]), val)
+
+        return logits
+
+    def _masked_softmax(self, logits: torch.Tensor, mask: Optional[AttentionMask]) -> torch.Tensor:
+        if mask is None or (mask.src_length_mask is None and mask.position_mask is None):
+            return F.softmax(logits, -1)
+
+        # Output shape: [n_batch * n_heads, n_time_dest, n_time_src]
+        bb, n_time_dest, n_time_src = logits.shape
+
+        logits = logits.view(bb // self.n_heads, self.n_heads, n_time_dest, n_time_src)
+        logits = self.apply_logit_masks(logits, mask)
+
+        logits = F.softmax(logits, -1)
+        return logits.view(bb, n_time_dest, n_time_src)
+
+    def _attention_read(self, mask: Optional[AttentionMask], scores: torch.Tensor, v: torch.Tensor) -> \
+            Tuple[torch.Tensor, torch.Tensor]:
+        # scores: [n_batch * n_heads, n_out, n_in]
+        # v: [n_nbatch * n_heads, n_in]
+        # Output data shape [n_batch * n_heads, n_time_dest, data_size]
+        # Out attention score shape: [n_batch, n_heads, n_time_dest, n_time_src]
+        s_reshape = scores.view(-1, self.n_heads, *scores.shape[1:])
+        # scores = self.dropout(scores)
+        if self.visualization_enabled:
+            self.attention_to_visualize.append(s_reshape[0])
+        return torch.bmm(scores, v), s_reshape
+
+    def transform_data(self, input: torch.Tensor, proj: Callable[[torch.Tensor], torch.Tensor],
+                       n_projs: int) -> List[torch.Tensor]:
+        # Input shape: [n_batch, n_steps, n_channels]
+        # Output: Tuple of n_projs tensors of dimension: [n_batch * n_heads, n_steps, projection_size]
+        n_batch, n_steps, _ = input.shape
+        transformed = proj(input).view(n_batch, n_steps, self.n_heads, n_projs, -1). \
+            permute(0, 2, 1, 3, 4).contiguous().view(n_batch * self.n_heads, n_steps, n_projs, -1)
+        return transformed.unbind(dim=2)
+
+    def plot(self, options: Dict[str, Any]) -> Dict[str, Any]:
+        r = {}
+        marks = options.get("steplabel")
+        if options.get("mha.plot_head_details") and self.attention_to_visualize[0].shape[0] > 1:
+            for head in range(self.attention_to_visualize[0].shape[0]):
+                r[f"head_{head}"] = framework.visualize.plot.AnimatedHeatmap(
+                        torch.stack([layer[head] for _, layer in enumerate(self.attention_to_visualize)], 0),
+                        ylabel="dest", xlabel="src", textval=False, x_marks=marks, y_marks=marks, ignore_wrong_marks=True)
+
+        r["attention_max"] = framework.visualize.plot.AnimatedHeatmap(
+            torch.stack([layer.max(0)[0] for _, layer in enumerate(self.attention_to_visualize)], 0),
+            ylabel="dest", xlabel="src", textval=False, x_marks=marks, y_marks=marks, ignore_wrong_marks=True)
+        self.attention_to_visualize = []
+        return r
+
+
+class AttentionMergeMixin:
+    def __init__(self, out_size: Optional[int]) -> None:
+        self.multi_head_merge = torch.nn.Linear(self.n_heads * self.projection_size, out_size or self.state_size, 
+                                       bias=False)
+
+    def merged_attention(self, n_batch: int, n_out_steps: int, *args, need_weights: bool = False, **kwargs) -> \
+            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+
+        data, scores = self._attention(*args, **kwargs)
+
+        data = data.view(n_batch, self.n_heads, n_out_steps, -1).permute(0, 2, 1, 3).contiguous().\
+                    view(n_batch, n_out_steps, -1)
+
+        return self.multi_head_merge(data), scores
+
+
+class AbsPosAttentionBase(MultiHeadAttentionBase):
+    def get_attention_scores(self, mask: Optional[torch.Tensor], q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
+        logits = torch.bmm(q, k.transpose(1, 2))
+        return self._masked_softmax(logits * self.scale, mask)
+
+    def _attention(self, mask: Optional[torch.Tensor], q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> \
+                   torch.Tensor:
+        # all inputs should have a shape of [n_batch, n_steps, data_size]
+        # Output shape [n_batch * n_heads, n_time_dest, data_size]
+        scores = self.get_attention_scores(mask, q, k)
+        return self._attention_read(mask, scores, v)
+
+
+class MultiHeadAttention(AttentionMergeMixin, AbsPosAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.1, input_size: Optional[int] = None,
+                 out_size: Optional[int] = None):
+        super(AttentionMergeMixin, self).__init__(state_size, n_heads, dropout)
+
+        self.data_to_kv = torch.nn.Linear(state_size, 2 * n_heads * self.projection_size, bias=False)
+        self.data_to_q = torch.nn.Linear(input_size or state_size, n_heads * self.projection_size, bias=False)
+
+        super(MultiHeadAttention, self).__init__(out_size)
+        self.reset_parameters()
+
+    def forward(self, curr_state: torch.Tensor, attend_to: torch.Tensor, mask: Optional[AttentionMask],
+                need_weights: bool = False):
+        # Input and output shape: [n_batch, n_steps, data_size]
+        k, v = self.transform_data(attend_to, self.data_to_kv, 2)
+        q, = self.transform_data(curr_state, self.data_to_q, 1)
+
+        data, scores = self.merged_attention(curr_state.shape[0], q.shape[1], mask, q, k, v)
+        if need_weights:
+            return data, scores
+        else:
+            return data
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.data_to_q.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)

ops/multi_head_relative_pos_attention.py

@@ -0,0 +1,185 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+from typing import Optional, Dict, Any
+from forgetting_transformer.ops.multi_head_attention import AttentionMask, MultiHeadAttentionBase, AttentionMergeMixin
+import forgetting_transformer.ops.framework_mock as framework
+import math
+from matplotlib import cm
+
+def shift(posmat: torch.Tensor) -> torch.Tensor:
+    # Slice out a matrix diagonally. Each successive row is sliced one position to the left compared.
+    # shape: [n_batch, n_head, n_out, n_in * 2 - 1]
+    # return: [n_batch, n_head, n_out, n_in]
+    p = F.pad(posmat, (0, 1, 0, 1)).flatten(-2)  # [n_batch, n_head, (n_out + 1) * n_in * 2]
+    p = p.narrow(-1, posmat.shape[-1] // 2, posmat.shape[-1] * posmat.shape[-2]).view_as(posmat)
+
+    return p.narrow(-1, 0, (posmat.shape[-1] + 1) // 2)
+
+
+class RelativeAttentionBase(MultiHeadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float, projection_size: Optional[int] = None):
+        super().__init__(state_size, n_heads, dropout=dropout, projection_size=projection_size)
+        self.scale = torch.nn.Parameter(torch.tensor([self.scale]))
+        self.s_bias = torch.nn.Parameter(torch.tensor([0.0]))
+        self.vis_pos_vs_content = []
+
+    def get_attention_scores(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor, k_pos: torch.Tensor, 
+                   pos_offset: int, ar_gate: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        # shape of q_content, q_pos, k_pos: [n_batch * n_heads, n_steps, data_size]
+        # k_pos: [n_heads, n_in * 2 - 1, data_size]
+        # ar_gate: [n_batch*n_heads, n_out, 1]
+        # Output shape [n_batch * n_heads, n_out, data_size]
+
+        n_batch = q_content.shape[0] // self.n_heads
+        n_out_steps = q_content.shape[1]
+
+        # content-content addressing
+        content = torch.bmm(q_content, self.dropout(k_content).transpose(1, 2))
+
+        # content-pos addressing.
+        pos = torch.matmul(q_pos.view(n_batch, self.n_heads, n_out_steps, -1), self.dropout(k_pos).transpose(-1, -2))  # [n_batch, n_head, n_out, n_in * 2 - 1]
+        fpos = shift(pos).flatten(0, 1)
+        if ar_gate is not None:
+            fpos = fpos * ar_gate + pos.flatten(0, 1)[..., fpos.shape[-1] - 1:] * (1 - ar_gate)
+        
+        # return self._masked_softmax((fpos) * self.scale, mask)
+        if self.visualization_enabled:
+            self.vis_pos_vs_content.append((content.view(n_batch, self.n_heads, *content.shape[1:])[0] * self.scale,
+                                            fpos.view(n_batch, self.n_heads, *fpos.shape[1:])[0] * self.scale))
+
+        return self._masked_softmax((content + fpos) * self.scale, mask)
+
+    def _attention(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor, k_pos: torch.Tensor,
+                   v: torch.Tensor, pos_offset: int,
+                   ar_gate: Optional[torch.Tensor] = None) -> [torch.Tensor, torch.Tensor]:
+
+        scores = self.get_attention_scores(mask, q_content, k_content, q_pos, k_pos, pos_offset, ar_gate)
+
+        # Scores shape: [n_batch * n_heads, n_out, n_in]
+        return self._attention_read(mask, scores, v)
+
+    def _get_pos_subset(self, pos_encoding: torch.Tensor, length: int, offset: int) -> torch.Tensor:
+        l_slice = 2 * length - 1
+        assert pos_encoding.shape[0] > l_slice
+        return pos_encoding.narrow(0, pos_encoding.shape[0] // 2 - length + 1 - offset, 2 * length - 1)
+
+    def plot(self, options: Dict[str, Any]) -> Dict[str, Any]:
+        r = {}
+        marks = options.get("steplabel")
+        if options.get("mha.plot_head_details") and self.vis_pos_vs_content:
+            for head in range(self.vis_pos_vs_content[0][0].shape[0]):
+                cont = torch.stack([layer[0][head] for _, layer in enumerate(self.vis_pos_vs_content)], 0)
+                pos = torch.stack([layer[1][head] for _, layer in enumerate(self.vis_pos_vs_content)], 0)
+                i = torch.stack([layer[head] for _, layer in enumerate(self.attention_to_visualize)], 0)
+                content = torch.stack([cont, pos], -1).softmax(-1)[..., 0]
+
+                color = cm.get_cmap("brg")(content.cpu().numpy())
+                color[..., -1] = (i * 0.95 + 0.05).cpu().numpy()
+
+                r[f"content_vs_pos_{head}"] = framework.visualize.plot.AnimatedHeatmap(color, ylabel="dest", 
+                    xlabel="src", textval=False, x_marks=marks, y_marks=marks, cmap="brg", colorbar=True,
+                    colorbar_ticks=[0, 0.99], colorbar_labels=["pos", "con"], ignore_wrong_marks=True)
+
+        # r["attention_max"] = framework.visualize.plot.AnimatedHeatmap(
+        #     torch.stack([layer.max(0)[0] for _, layer in enumerate(self.attention_to_visualize)], 0),
+        #     ylabel="dest", xlabel="src", textval=False, x_marks=marks, y_marks=marks)
+        self.vis_pos_vs_content = []
+
+        r.update(super().plot(options))
+        return r
+
+
+
+class FixedRelativeMultiheadAttentionBase(RelativeAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, input_size: Optional[int] = None,
+                 projection_size: Optional[int] = None):
+        super().__init__(state_size, n_heads, dropout, projection_size)
+
+        self.input_size = state_size if input_size is None else input_size
+
+        self.pos_to_pq = torch.nn.Linear(state_size, self.n_heads * self.projection_size, bias=False)
+        self.register_buffer("pos_encoding", self._create_buffer(1000))
+
+    def _create_buffer(self, max_len: int):
+        return framework.layers.sinusoidal_pos_embedding(self.state_size, 2 * max_len - 1, -max_len + 1,
+                                                         device=self.pos_to_pq.weight.device)
+
+    def get_pos(self, l: int, offset: int) -> torch.Tensor:
+        if self.pos_encoding.shape[0] < 2 * (l + offset) - 1:
+            self.pos_encoding = self._create_buffer(int(2**math.ceil(math.log2(2 * (l + offset) - 1))))
+
+        return self.pos_to_pq(self._get_pos_subset(self.pos_encoding, l, offset))
+
+
+class FixedRelativeMultiheadAttention(AttentionMergeMixin, FixedRelativeMultiheadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, global_pos_bias: bool = True,
+                 global_content_bias: bool = True, input_size: Optional[int] = None, absolute_gate: bool = False,
+                 projection_size: Optional[int] = None, output_size: Optional[int] = None):
+        super(AttentionMergeMixin, self).__init__(state_size, n_heads, dropout, input_size, projection_size=projection_size)
+
+        self.data_to_kv = torch.nn.Linear(state_size, 2 * n_heads * self.projection_size, bias=False)
+        self.data_to_q = torch.nn.Linear(self.input_size, n_heads * self.projection_size, bias=False)
+        self.data_to_absgate = torch.nn.Linear(self.input_size, n_heads) \
+                               if absolute_gate else None
+
+        self.global_content_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                                   if global_content_bias else None
+        self.global_pos_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                               if global_pos_bias else None
+
+        super(FixedRelativeMultiheadAttention, self).__init__(output_size)
+        self.reset_parameters()
+
+    def add_head_specific_bias(self, data: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor:
+        # data [batch * n_heads, len, c]
+        # bias [n_heads, c]
+        return (data.view(-1, bias.shape[0], *data.shape[1:]) + bias.unsqueeze(1).type_as(data)).view_as(data) \
+               if bias is not None else data
+
+    def forward(self, curr_state: torch.Tensor, attend_to: torch.Tensor, mask: Optional[AttentionMask],
+                pos_offset: int = 0, need_weights: bool = False):
+        # curr_state: [§size, out_len, c]
+        # attend_to: [batch_size, in_len, c]
+        batch_size, in_len = attend_to.shape[0:2]
+        out_len = curr_state.shape[1]
+
+        k_content, v = self.transform_data(attend_to, self.data_to_kv, 2)
+        q, = self.transform_data(curr_state, self.data_to_q, 1)
+
+        k_pos = self.get_pos(in_len, pos_offset).view(-1, self.n_heads, self.projection_size).\
+                transpose(0, 1)  # n_heads, 2*in_len -1 , projection_size
+
+        q_content = self.add_head_specific_bias(q, self.global_content_bias)
+        q_pos = self.add_head_specific_bias(q, self.global_pos_bias)
+
+        
+        absgate = torch.sigmoid(self.transform_data(curr_state, self.data_to_absgate, 1)[0]) \
+                  if self.data_to_absgate is not None else None
+    
+        data, scores = self.merged_attention(batch_size, out_len, mask, q_content, k_content, q_pos, k_pos, v, 
+                                             pos_offset, ar_gate=absgate, need_weights=need_weights)
+
+        if need_weights:
+            return data, scores
+        else:
+            return data
+
+    def reset_parameters(self):
+    #     # super().reset_parameters()
+
+        torch.nn.init.xavier_uniform_(self.data_to_q.weight)
+        torch.nn.init.xavier_uniform_(self.pos_to_pq.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)
+
+        if self.global_pos_bias is not None:
+            self.global_pos_bias.data.fill_(0)

ops/multi_head_relative_pos_attention.py.bak

@@ -0,0 +1,185 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+from typing import Optional, Dict, Any
+from .multi_head_attention import AttentionMask, MultiHeadAttentionBase, AttentionMergeMixin
+import framework
+import math
+from matplotlib import cm
+
+def shift(posmat: torch.Tensor) -> torch.Tensor:
+    # Slice out a matrix diagonally. Each successive row is sliced one position to the left compared.
+    # shape: [n_batch, n_head, n_out, n_in * 2 - 1]
+    # return: [n_batch, n_head, n_out, n_in]
+    p = F.pad(posmat, (0, 1, 0, 1)).flatten(-2)  # [n_batch, n_head, (n_out + 1) * n_in * 2]
+    p = p.narrow(-1, posmat.shape[-1] // 2, posmat.shape[-1] * posmat.shape[-2]).view_as(posmat)
+
+    return p.narrow(-1, 0, (posmat.shape[-1] + 1) // 2)
+
+
+class RelativeAttentionBase(MultiHeadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float, projection_size: Optional[int] = None):
+        super().__init__(state_size, n_heads, dropout=dropout, projection_size=projection_size)
+        self.scale = torch.nn.Parameter(torch.tensor([self.scale]))
+        self.s_bias = torch.nn.Parameter(torch.tensor([0.0]))
+        self.vis_pos_vs_content = []
+
+    def get_attention_scores(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor, k_pos: torch.Tensor, 
+                   pos_offset: int, ar_gate: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        # shape of q_content, q_pos, k_pos: [n_batch * n_heads, n_steps, data_size]
+        # k_pos: [n_heads, n_in * 2 - 1, data_size]
+        # ar_gate: [n_batch*n_heads, n_out, 1]
+        # Output shape [n_batch * n_heads, n_out, data_size]
+
+        n_batch = q_content.shape[0] // self.n_heads
+        n_out_steps = q_content.shape[1]
+
+        # content-content addressing
+        content = torch.bmm(q_content, self.dropout(k_content).transpose(1, 2))
+
+        # content-pos addressing.
+        pos = torch.matmul(q_pos.view(n_batch, self.n_heads, n_out_steps, -1), self.dropout(k_pos).transpose(-1, -2))  # [n_batch, n_head, n_out, n_in * 2 - 1]
+        fpos = shift(pos).flatten(0, 1)
+        if ar_gate is not None:
+            fpos = fpos * ar_gate + pos.flatten(0, 1)[..., fpos.shape[-1] - 1:] * (1 - ar_gate)
+        
+        # return self._masked_softmax((fpos) * self.scale, mask)
+        if self.visualization_enabled:
+            self.vis_pos_vs_content.append((content.view(n_batch, self.n_heads, *content.shape[1:])[0] * self.scale,
+                                            fpos.view(n_batch, self.n_heads, *fpos.shape[1:])[0] * self.scale))
+
+        return self._masked_softmax((content + fpos) * self.scale, mask)
+
+    def _attention(self, mask: Optional[torch.Tensor],
+                   q_content: torch.Tensor, k_content: torch.Tensor,
+                   q_pos: torch.Tensor, k_pos: torch.Tensor,
+                   v: torch.Tensor, pos_offset: int,
+                   ar_gate: Optional[torch.Tensor] = None) -> [torch.Tensor, torch.Tensor]:
+
+        scores = self.get_attention_scores(mask, q_content, k_content, q_pos, k_pos, pos_offset, ar_gate)
+
+        # Scores shape: [n_batch * n_heads, n_out, n_in]
+        return self._attention_read(mask, scores, v)
+
+    def _get_pos_subset(self, pos_encoding: torch.Tensor, length: int, offset: int) -> torch.Tensor:
+        l_slice = 2 * length - 1
+        assert pos_encoding.shape[0] > l_slice
+        return pos_encoding.narrow(0, pos_encoding.shape[0] // 2 - length + 1 - offset, 2 * length - 1)
+
+    def plot(self, options: Dict[str, Any]) -> Dict[str, Any]:
+        r = {}
+        marks = options.get("steplabel")
+        if options.get("mha.plot_head_details") and self.vis_pos_vs_content:
+            for head in range(self.vis_pos_vs_content[0][0].shape[0]):
+                cont = torch.stack([layer[0][head] for _, layer in enumerate(self.vis_pos_vs_content)], 0)
+                pos = torch.stack([layer[1][head] for _, layer in enumerate(self.vis_pos_vs_content)], 0)
+                i = torch.stack([layer[head] for _, layer in enumerate(self.attention_to_visualize)], 0)
+                content = torch.stack([cont, pos], -1).softmax(-1)[..., 0]
+
+                color = cm.get_cmap("brg")(content.cpu().numpy())
+                color[..., -1] = (i * 0.95 + 0.05).cpu().numpy()
+
+                r[f"content_vs_pos_{head}"] = framework.visualize.plot.AnimatedHeatmap(color, ylabel="dest", 
+                    xlabel="src", textval=False, x_marks=marks, y_marks=marks, cmap="brg", colorbar=True,
+                    colorbar_ticks=[0, 0.99], colorbar_labels=["pos", "con"], ignore_wrong_marks=True)
+
+        # r["attention_max"] = framework.visualize.plot.AnimatedHeatmap(
+        #     torch.stack([layer.max(0)[0] for _, layer in enumerate(self.attention_to_visualize)], 0),
+        #     ylabel="dest", xlabel="src", textval=False, x_marks=marks, y_marks=marks)
+        self.vis_pos_vs_content = []
+
+        r.update(super().plot(options))
+        return r
+
+
+
+class FixedRelativeMultiheadAttentionBase(RelativeAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, input_size: Optional[int] = None,
+                 projection_size: Optional[int] = None):
+        super().__init__(state_size, n_heads, dropout, projection_size)
+
+        self.input_size = state_size if input_size is None else input_size
+
+        self.pos_to_pq = torch.nn.Linear(state_size, self.n_heads * self.projection_size, bias=False)
+        self.register_buffer("pos_encoding", self._create_buffer(1000))
+
+    def _create_buffer(self, max_len: int):
+        return framework.layers.sinusoidal_pos_embedding(self.state_size, 2 * max_len - 1, -max_len + 1,
+                                                         device=self.pos_to_pq.weight.device)
+
+    def get_pos(self, l: int, offset: int) -> torch.Tensor:
+        if self.pos_encoding.shape[0] < 2 * (l + offset) - 1:
+            self.pos_encoding = self._create_buffer(int(2**math.ceil(math.log2(2 * (l + offset) - 1))))
+
+        return self.pos_to_pq(self._get_pos_subset(self.pos_encoding, l, offset))
+
+
+class FixedRelativeMultiheadAttention(AttentionMergeMixin, FixedRelativeMultiheadAttentionBase):
+    def __init__(self, state_size: int, n_heads: int, dropout: float = 0.0, global_pos_bias: bool = True,
+                 global_content_bias: bool = True, input_size: Optional[int] = None, absolute_gate: bool = False,
+                 projection_size: Optional[int] = None, output_size: Optional[int] = None):
+        super(AttentionMergeMixin, self).__init__(state_size, n_heads, dropout, input_size, projection_size=projection_size)
+
+        self.data_to_kv = torch.nn.Linear(state_size, 2 * n_heads * self.projection_size, bias=False)
+        self.data_to_q = torch.nn.Linear(self.input_size, n_heads * self.projection_size, bias=False)
+        self.data_to_absgate = torch.nn.Linear(self.input_size, n_heads) \
+                               if absolute_gate else None
+
+        self.global_content_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                                   if global_content_bias else None
+        self.global_pos_bias = torch.nn.Parameter(torch.zeros([n_heads, self.projection_size])) \
+                               if global_pos_bias else None
+
+        super(FixedRelativeMultiheadAttention, self).__init__(output_size)
+        self.reset_parameters()
+
+    def add_head_specific_bias(self, data: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor:
+        # data [batch * n_heads, len, c]
+        # bias [n_heads, c]
+        return (data.view(-1, bias.shape[0], *data.shape[1:]) + bias.unsqueeze(1).type_as(data)).view_as(data) \
+               if bias is not None else data
+
+    def forward(self, curr_state: torch.Tensor, attend_to: torch.Tensor, mask: Optional[AttentionMask],
+                pos_offset: int = 0, need_weights: bool = False):
+        # curr_state: [§size, out_len, c]
+        # attend_to: [batch_size, in_len, c]
+        batch_size, in_len = attend_to.shape[0:2]
+        out_len = curr_state.shape[1]
+
+        k_content, v = self.transform_data(attend_to, self.data_to_kv, 2)
+        q, = self.transform_data(curr_state, self.data_to_q, 1)
+
+        k_pos = self.get_pos(in_len, pos_offset).view(-1, self.n_heads, self.projection_size).\
+                transpose(0, 1)  # n_heads, 2*in_len -1 , projection_size
+
+        q_content = self.add_head_specific_bias(q, self.global_content_bias)
+        q_pos = self.add_head_specific_bias(q, self.global_pos_bias)
+
+        
+        absgate = torch.sigmoid(self.transform_data(curr_state, self.data_to_absgate, 1)[0]) \
+                  if self.data_to_absgate is not None else None
+    
+        data, scores = self.merged_attention(batch_size, out_len, mask, q_content, k_content, q_pos, k_pos, v, 
+                                             pos_offset, ar_gate=absgate, need_weights=need_weights)
+
+        if need_weights:
+            return data, scores
+        else:
+            return data
+
+    def reset_parameters(self):
+    #     # super().reset_parameters()
+
+        torch.nn.init.xavier_uniform_(self.data_to_q.weight)
+        torch.nn.init.xavier_uniform_(self.pos_to_pq.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)
+        torch.nn.init.xavier_uniform_(self.data_to_kv.weight)
+
+        if self.global_content_bias is not None:
+            self.global_content_bias.data.fill_(0)
+
+        if self.global_pos_bias is not None:
+            self.global_pos_bias.data.fill_(0)

ops/sliding_window_attention_std.py

@@ -0,0 +1,88 @@
+"""
+Sliding Window / Hard Attention
+Based on "Context Limitations Make Neural Language Models More Human-Like"
+(Kuribayashi et al., 2022)
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+def sliding_window_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    window_size: int = 2,  # 默认2-gram（看前1个token）
+) -> torch.Tensor:
+    """
+    Sliding Window Attention
+    
+    硬截断：只能attend到最近window_size个token
+    """
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    T_k = k.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # Compute logits
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # Create sliding window mask
+    mask = create_sliding_window_mask(T_q, T_k, window_size, device=q.device)
+    logits = logits.masked_fill(~mask, float('-inf'))
+    
+    # Seq start mask
+    if seq_start is not None:
+        seq_mask = torch.arange(T_k, device=q.device)[None, None, None, :] < seq_start[None, :, None, None]
+        logits = logits.masked_fill(seq_mask, float('-inf'))
+    
+    # Standard softmax
+    weights = F.softmax(logits, dim=-1)
+    
+    # Apply to values
+    out = torch.matmul(weights, v)
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out
+
+
+def create_sliding_window_mask(
+    T_q: int,
+    T_k: int,
+    window_size: int,
+    device: torch.device
+) -> torch.Tensor:
+    """
+    创建sliding window mask
+    
+    window_size=1: 只看前1个token (2-gram)
+    window_size=2: 只看前2个token (3-gram)
+    """
+    # 基础causal mask
+    mask = torch.tril(torch.ones(T_q, T_k, dtype=torch.bool, device=device))
+    
+    # 应用window限制
+    if window_size > 0 and window_size < T_k:
+        for i in range(T_q):
+            # 只保留 [i-window_size+1, i] 范围
+            start = max(0, i - window_size + 1)
+            if start > 0:
+                mask[i, :start] = False
+    
+    return mask[None, None, :, :]  # [1, 1, T_q, T_k]

ops/stickbreaking_attention_std.py

@@ -0,0 +1,46 @@
+"""
+Stick-breaking Attention - 官方Triton实现
+"""
+
+from stickbreaking_attention.sb_attn import sb_attn
+import math
+import torch
+from einops import rearrange
+from typing import Optional
+
+
+def stickbreaking_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    head_first: bool = False,
+    seq_start: Optional[torch.Tensor] = None,
+    sm_scale: Optional[float] = None,
+    normalize: bool = True,
+    attend_current: bool = False,
+) -> torch.Tensor:
+    """Stick-breaking attention using official Triton implementation"""
+    
+    if not head_first:
+        q = rearrange(q, "b t h d -> b h t d")
+        k = rearrange(k, "b t h d -> b h t d")
+        v = rearrange(v, "b t h d -> b h t d")
+    
+    B, H, T_q, D = q.shape
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # 官方Triton实现
+    # 返回 (output, remainder)
+    out, rem = sb_attn(
+        q, k, v, 
+        inv_temp=sm_scale,
+        attend_current=attend_current
+    )
+    
+    if not head_first:
+        out = rearrange(out, "b h t d -> b t h d")
+    
+    return out

ops/transformer.py

@@ -0,0 +1,165 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+from .multi_head_attention import MultiHeadAttention, AttentionMask
+from typing import Optional, Callable, Dict
+from dataclasses import dataclass
+# This file is based on PyTorch's internal implementation
+
+ActivationFunction = Callable[[torch.Tensor], torch.Tensor]
+
+
+class TransformerEncoderLayer(torch.nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation: ActivationFunction = F.relu,
+                 attention_dropout=0):
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout)
+        self.linear1 = torch.nn.Linear(d_model, dim_feedforward)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.linear2 = torch.nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = torch.nn.LayerNorm(d_model)
+        self.norm2 = torch.nn.LayerNorm(d_model)
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+
+        self.activation = activation
+        self.reset_parameters()
+
+    def forward(self, src: torch.Tensor, mask: Optional[AttentionMask] = None) -> torch.Tensor:
+        src2 = self.self_attn(src, src, mask)
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.linear1.weight, gain=torch.nn.init.calculate_gain('relu') \
+            if self.activation is F.relu else 1.0)
+        torch.nn.init.xavier_uniform_(self.linear2.weight)
+
+
+class TransformerDecoderLayer(torch.nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation: ActivationFunction = F.relu,
+                 attention_dropout=0):
+        super(TransformerDecoderLayer, self).__init__()
+
+        self.self_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout)
+        self.multihead_attn = MultiHeadAttention(d_model, nhead, dropout=attention_dropout)
+        # Implementation of Feedforward model
+        self.linear1 = torch.nn.Linear(d_model, dim_feedforward)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.linear2 = torch.nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = torch.nn.LayerNorm(d_model)
+        self.norm2 = torch.nn.LayerNorm(d_model)
+        self.norm3 = torch.nn.LayerNorm(d_model)
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+        self.dropout3 = torch.nn.Dropout(dropout)
+
+        self.activation = activation
+        self.reset_parameters()
+
+    def forward(self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None,
+                memory_key_padding_mask: Optional[torch.Tensor] = None,
+                full_target: Optional[torch.Tensor] = None, pos_offset: int = 0) -> torch.Tensor:
+        
+        assert pos_offset == 0 or tgt_mask is None
+        tgt2 = self.self_attn(tgt, tgt if full_target is None else full_target, mask=AttentionMask(None, tgt_mask))
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(tgt, memory, mask=AttentionMask(memory_key_padding_mask, None))
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.linear1.weight, gain=torch.nn.init.calculate_gain('relu') \
+            if self.activation is F.relu else 1.0)
+        torch.nn.init.xavier_uniform_(self.linear2.weight)
+
+
+class TransformerDecoderBase(torch.nn.Module):
+    @dataclass
+    class State:
+        step: int
+        state: Dict[int, torch.Tensor]
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+    def create_state(self, batch_size: int, max_length: int, device: torch.device) -> State:
+        return self.State(0, {i: torch.empty([batch_size, max_length, self.d_model], device=device)
+                              for i in range(len(self.layers))})
+
+    def one_step_forward(self, state: State, data: torch.Tensor, *args, **kwargs):
+        assert data.shape[1] == 1, f"For one-step forward should have one timesteps, but shape is {data.shape}"
+        assert state.step < state.state[0].shape[1]
+
+        for i, l in enumerate(self.layers):
+            state.state[i][:, state.step:state.step + 1] = data
+            data = l(data, *args, **kwargs, full_target=state.state[i][:, :state.step + 1],
+                     pos_offset=state.step)
+
+        state.step += 1
+        return data
+
+
+class TransformerEncoder(torch.nn.Module):
+    def __init__(self, layer, n_layers: int, *args, **kwargs):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([layer(*args, **kwargs) for _ in range(n_layers)])
+
+    def forward(self, data: torch.Tensor, *args, **kwargs):
+        for l in self.layers:
+            data = l(data, *args, **kwargs)
+        return data
+
+
+class TransformerDecoder(TransformerDecoderBase):
+    def __init__(self, layer, n_layers: int, d_model: int, *args, **kwargs):
+        super().__init__(d_model)
+        self.layers = torch.nn.ModuleList([layer(d_model, *args, **kwargs) for _ in range(n_layers)])
+
+    def forward(self, data: torch.Tensor, *args, **kwargs):
+        for l in self.layers:
+            data = l(data, *args, **kwargs)
+        return data
+
+
+def TransformerEncoderWithLayer(layer = TransformerEncoder):
+    return lambda *args, **kwargs: TransformerEncoder(layer, *args, **kwargs)
+
+
+def TransformerDecoderWithLayer(layer = TransformerDecoder):
+    return lambda *args, **kwargs: TransformerDecoder(layer, *args, **kwargs)
+
+
+class Transformer(torch.nn.Module):
+    def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
+                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
+                 activation: ActivationFunction = F.relu, encoder_layer=TransformerEncoderWithLayer(),
+                 decoder_layer=TransformerDecoderWithLayer(), attention_dropout: float = 0):
+        super().__init__()
+
+        self.encoder = encoder_layer(num_encoder_layers, d_model, nhead, dim_feedforward,
+                                     dropout, activation, attention_dropout)
+        self.decoder = decoder_layer(num_decoder_layers, d_model, nhead, dim_feedforward,
+                                     dropout, activation, attention_dropout)
+
+    def forward(self, src: torch.Tensor, tgt: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None,
+                src_mask: Optional[AttentionMask] = None):
+
+        memory = self.encoder(src, src_mask)
+        return self.decoder(tgt, memory, tgt_mask, src_mask.src_length_mask if src_mask is not None else None)
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz: int, device: torch.device) -> torch.Tensor:
+        return torch.triu(torch.ones(sz, sz, dtype=torch.bool, device=device), diagonal=1)

ops/vanilla_attention_std.py

@@ -0,0 +1,171 @@
+"""
+Vanilla Transformer 的标准 Softmax Attention
+用于替换 flash_attn 的实现
+"""
+import math
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional, Tuple
+
+def vanilla_attention_std(
+    q: torch.Tensor,
+    k: torch.Tensor, 
+    v: torch.Tensor,
+    causal: bool = True,
+    window_size: Optional[Tuple[int, int]] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """
+    标准 Softmax Attention，兼容 flash_attn_func 的输入格式
+    
+    Args:
+        q, k, v: [batch, seq_len, num_heads, head_dim] 格式
+        causal: 是否使用因果mask
+        window_size: 滑动窗口大小 (left, right)，(-1, -1) 表示无限制
+        sm_scale: softmax 缩放因子
+    
+    Returns:
+        output: [batch, seq_len, num_heads, head_dim] 格式
+    """
+    B, T_q, H, D = q.shape
+    T_k = k.shape[1]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    # 转换为 [B, H, T, D] 格式进行计算
+    q = rearrange(q, 'b t h d -> b h t d')
+    k = rearrange(k, 'b t h d -> b h t d')
+    v = rearrange(v, 'b t h d -> b h t d')
+    
+    # 计算 attention scores
+    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale
+    
+    # Causal mask
+    if causal:
+        P_SEQ = T_k - T_q  # 处理 KV cache 的情况
+        causal_mask = torch.triu(
+            torch.ones((T_q, T_k), dtype=torch.bool, device=q.device),
+            diagonal=P_SEQ + 1
+        )
+        scores = scores.masked_fill(causal_mask[None, None, :, :], float('-inf'))
+    
+    # Window mask (sliding window attention)
+    if window_size is not None and window_size != (-1, -1):
+        left_window, right_window = window_size
+        window_mask = torch.ones((T_q, T_k), dtype=torch.bool, device=q.device)
+        for i in range(T_q):
+            # 计算每个查询位置的有效窗口范围
+            start = max(0, i - left_window)
+            end = min(T_k, i + right_window + 1)
+            window_mask[i, start:end] = False
+        scores = scores.masked_fill(window_mask[None, None, :, :], float('-inf'))
+    
+    # Softmax
+    attn_weights = F.softmax(scores, dim=-1)
+    attn_weights = torch.nan_to_num(attn_weights, 0.0)
+    
+    # Apply attention to values
+    output = torch.matmul(attn_weights.to(v.dtype), v)
+    
+    # 转换回 [B, T, H, D] 格式
+    output = rearrange(output, 'b h t d -> b t h d')
+    
+    return output
+
+
+def vanilla_attention_varlen_std(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    causal: bool = True,
+    window_size: Optional[Tuple[int, int]] = None,
+    sm_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """
+    变长序列的标准 Softmax Attention，兼容 flash_attn_varlen_func
+    
+    Args:
+        q: [total_q_tokens, num_heads, head_dim]
+        k: [total_k_tokens, num_kv_heads, head_dim]
+        v: [total_k_tokens, num_kv_heads, head_dim]
+        cu_seqlens_q: 累积序列长度 [batch_size + 1]
+        cu_seqlens_k: 累积序列长度 [batch_size + 1]
+        max_seqlen_q: 最大查询序列长度
+        max_seqlen_k: 最大键值序列长度
+    
+    Returns:
+        output: [total_q_tokens, num_heads, head_dim]
+    """
+    batch_size = cu_seqlens_q.shape[0] - 1
+    H = q.shape[1]
+    D = q.shape[2]
+    
+    if sm_scale is None:
+        sm_scale = 1.0 / math.sqrt(D)
+    
+    outputs = []
+    
+    # 逐批次处理
+    for b in range(batch_size):
+        q_start, q_end = cu_seqlens_q[b].item(), cu_seqlens_q[b+1].item()
+        k_start, k_end = cu_seqlens_k[b].item(), cu_seqlens_k[b+1].item()
+        
+        if q_start == q_end:  # 空序列
+            continue
+            
+        # 提取当前批次的 q, k, v
+        q_b = q[q_start:q_end]  # [T_q, H, D]
+        k_b = k[k_start:k_end]  # [T_k, H, D] 
+        v_b = v[k_start:k_end]  # [T_k, H, D]
+        
+        T_q = q_b.shape[0]
+        T_k = k_b.shape[0]
+        
+        # 转换为 [H, T, D] 格式
+        q_b = rearrange(q_b, 't h d -> h t d')
+        k_b = rearrange(k_b, 't h d -> h t d')
+        v_b = rearrange(v_b, 't h d -> h t d')
+        
+        # 计算 attention scores
+        scores = torch.matmul(q_b.float(), k_b.float().transpose(-2, -1)) * sm_scale
+        
+        # Causal mask
+        if causal:
+            P_SEQ = T_k - T_q
+            causal_mask = torch.triu(
+                torch.ones((T_q, T_k), dtype=torch.bool, device=q.device),
+                diagonal=P_SEQ + 1
+            )
+            scores = scores.masked_fill(causal_mask[None, :, :], float('-inf'))
+        
+        # Window mask
+        if window_size is not None and window_size != (-1, -1):
+            left_window, right_window = window_size
+            window_mask = torch.ones((T_q, T_k), dtype=torch.bool, device=q.device)
+            for i in range(T_q):
+                start = max(0, i - left_window)
+                end = min(T_k, i + right_window + 1)
+                window_mask[i, start:end] = False
+            scores = scores.masked_fill(window_mask[None, :, :], float('-inf'))
+        
+        # Softmax
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = torch.nan_to_num(attn_weights, 0.0)
+        
+        # Apply attention
+        output_b = torch.matmul(attn_weights.to(v_b.dtype), v_b)
+        
+        # 转换回 [T, H, D] 格式
+        output_b = rearrange(output_b, 'h t d -> t h d')
+        outputs.append(output_b)
+    
+    # 拼接所有批次的输出
+    output = torch.cat(outputs, dim=0)
+    
+    return output