Update modeling_rwkv6qwen2.py

bugfix for FLA and transformers lib updates

modeling_rwkv6qwen2.py

@@ -51,6 +51,7 @@ from transformers.utils import (
 from .configuration_rwkv6qwen2 import RWKV6Qwen2Config
 
 from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer, Qwen2MLP, Qwen2RMSNorm, repeat_kv
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 
 logger = logging.get_logger(__name__)
 
@@ -58,7 +59,7 @@ logger = logging.get_logger(__name__)
 _CHECKPOINT_FOR_DOC = "RWKV/RWKV6Qwen2-7B"
 _CONFIG_FOR_DOC = "RWKV6Qwen2Config"
 
-class RWKV6State(Cache):
+class RWKV6State():
     def __init__(self) -> None:
         super().__init__()
         self._seen_tokens = 0  # Used in `generate` to keep tally of how many tokens the cache has seen
@@ -113,28 +114,6 @@ class RWKV6State(Cache):
         """
         return None
 
-    # def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
-    #     """Converts the `DynamicCache` instance into the its equivalent in the legacy cache format. Used for
-    #     backward compatibility."""
-    #     legacy_cache = ()
-    #     for layer_idx in range(len(self)):
-    #         legacy_cache += ((self.layer_kv_states[layer_idx], self.layer_shift_states[layer_idx]),)
-    #     return legacy_cache
-
-    # @classmethod
-    # #@deprecate_kwarg("num_hidden_layers", version="4.47.0")
-    # def from_legacy_cache(
-    #     cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor, torch.FloatTensor]]] = None, num_hidden_layers: int | None = None
-    # ) -> "RWKV6State":
-    #     """Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for
-    #     backward compatibility."""
-    #     cache = cls()
-    #     if past_key_values is not None:
-    #         for layer_idx in range(len(past_key_values)):
-    #             layer_kv_state, layer_shift_state = past_key_values[layer_idx]
-    #             cache.update(layer_kv_state, layer_shift_state, layer_idx)
-    #     return cache
-
     def crop(self, max_length: int):
         # can't implement this for linear attention variants
         return
@@ -144,8 +123,8 @@ class RWKV6State(Cache):
         self,
         kv_state: torch.Tensor,
         shift_state: torch.Tensor,
-        token_count: int,
         layer_idx: int,
+        token_count: int = 0,
         cache_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:        
         # Update the number of seen tokens
@@ -162,62 +141,140 @@ class RWKV6State(Cache):
 
         return self.layer_kv_states[layer_idx], self.layer_shift_states[layer_idx]
 
-    # @deprecate_kwarg("num_hidden_layers", version="4.47.0")
-    # def batch_split(
-    #     self, full_batch_size: int, split_size: int, num_hidden_layers: int = None
-    # ) -> List["DynamicCache"]:
-    #     """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by
-    #     `_split_model_inputs()` in `generation.utils`"""
-    #     out = []
-    #     for i in range(0, full_batch_size, split_size):
-    #         current_split = DynamicCache()
-    #         current_split._seen_tokens = self._seen_tokens
-    #         current_split.key_cache = [tensor[i : i + split_size] for tensor in self.key_cache]
-    #         current_split.value_cache = [tensor[i : i + split_size] for tensor in self.value_cache]
-    #         out.append(current_split)
-    #     return out
-
-    # @classmethod
-    # @deprecate_kwarg("num_hidden_layers", version="4.47.0")
-    # def from_batch_splits(cls, splits: List["DynamicCache"], num_hidden_layers: int = None) -> "DynamicCache":
-    #     """This is the opposite of the above `batch_split()` method. This will be used by `stack_model_outputs` in
-    #     `generation.utils`"""
-    #     cache = cls()
-    #     for idx in range(len(splits[0])):
-    #         key_cache = [current.key_cache[idx] for current in splits if current.key_cache[idx] != []]
-    #         value_cache = [current.key_cache[idx] for current in splits if current.key_cache[idx] != []]
-    #         if key_cache != []:
-    #             layer_keys = torch.cat(key_cache, dim=0)
-    #             layer_values = torch.cat(value_cache, dim=0)
-    #             cache.update(layer_keys, layer_values, idx)
-    #     return cache
-
-    # def batch_repeat_interleave(self, repeats: int):
-    #     """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search."""
-    #     for layer_idx in range(len(self)):
-    #         self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(repeats, dim=0)
-    #         self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(repeats, dim=0)
-
-    # def batch_select_indices(self, indices: torch.Tensor):
-    #     """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search."""
-    #     for layer_idx in range(len(self)):
-    #         self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...]
-    #         self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...]
-
 try:
     #from fla.ops.gla.chunk import chunk_gla
     from fla.ops.gla.fused_recurrent import fused_recurrent_gla
 except ImportError:
     print("Required module is not installed. Please install it using the following commands:")
-    print("pip install -U git+https://github.com/fla-org/flash-linear-attention")
+    print("pip install --no-use-pep517 flash-linear-attention")
     print("Additionally, ensure you have at least version 2.2.0 of Triton installed:")
     print("pip install triton>=2.2.0")
 
+class Qwen2RotaryEmbedding(nn.Module):
+    def __init__(self, config: RWKV6Qwen2Config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        """
+        dynamic RoPE layers should recompute `inv_freq` in the following situations:
+        1 - growing beyond the cached sequence length (allow scaling)
+        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
+        """
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
+            # This .to() is needed if the model has been moved to a device after being initialized (because
+            # the buffer is automatically moved, but not the original copy)
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+
+        # Core RoPE block
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
+        cos = cos * self.attention_scaling
+        sin = sin * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+    
+def generate_rotary_embedding(max_seqlen:int, dim:int, theta:float = 10000.0, scale:float = 1):
+    #inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float).to(device) / dim))
+
+    angular_velocity = theta ** -(torch.arange(0, dim, 2, dtype=torch.float) / dim) / scale # frequencies from 1.0 ... 1/theta
+    angles = torch.outer(torch.arange(max_seqlen), angular_velocity)
+    # Different from paper, but it uses a different permutation in order to obtain the same calculation
+    emb = torch.cat((angles, angles), dim=-1)
+    return torch.stack([emb.cos(), emb.sin()], dim=0)
+    #return torch.polar(torch.ones_like(angles), angles)
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def ortho_init(x, scale):
+    with torch.no_grad():
+        shape = x.shape
+        if len(shape) == 2:
+            gain = math.sqrt(shape[0] / shape[1]) if shape[0] > shape[1] else 1
+            #nn.init.orthogonal_(x, gain=gain * scale)
+            x.copy_(nn.init.orthogonal_(torch.empty_like(x, dtype=torch.float32), gain=gain * scale))
+        elif len(shape) == 3:
+            gain = math.sqrt(shape[1] / shape[2]) if shape[1] > shape[2] else 1
+            for i in range(shape[0]):
+                #nn.init.orthogonal_(x[i], gain=gain * scale)
+                x[i].copy_(nn.init.orthogonal_(torch.empty_like(x[i], dtype=torch.float32), gain=gain * scale))                            
+        else:
+            assert False
+        return x
+
 class RWKV6Attention(nn.Module):
     def __init__(self, config, layer_idx: Optional[int] = None):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
+
         if layer_idx is None:
             logger.warning_once(
                 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
@@ -232,6 +289,11 @@ class RWKV6Attention(nn.Module):
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.attention_dropout = config.attention_dropout
 
+        n_layer = self.config.num_hidden_layers
+        n_embd = self.hidden_size
+        dim_att = self.num_heads * self.head_dim
+        layer_id = self.layer_idx
+
         if self.hidden_size % self.num_heads != 0:
             raise ValueError(
                 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
@@ -242,41 +304,55 @@ class RWKV6Attention(nn.Module):
         self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=getattr(config, 'attention_output_bias', config.attention_bias))
 
-        self.gate = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        nn.init.zeros_(self.gate.weight)
+        calc_lora_rank = lambda exponent, multiplier: max(1, round(self.hidden_size ** exponent * multiplier / 32)) * 32
 
-        n_layer = self.config.num_hidden_layers
-        n_embd = self.hidden_size
-        dim_att = self.num_heads * self.head_dim
-        layer_id = self.layer_idx
+        if config.gate_rank_type == 1:
+            self.gate = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        elif config.gate_rank_type == 2:
+            lora_rank_gate = config.lora_rank_gate or calc_lora_rank(0.8, 0.6)
+            self.g1 = nn.Parameter(torch.empty(n_embd, lora_rank_gate))
+            self.g2 = nn.Parameter(torch.empty(lora_rank_gate, n_embd))
+
+        if config.groupnorm_att:
+            self.ln_x = nn.GroupNorm(self.num_heads, dim_att, eps=self.head_dim * 1e-5)
 
         with torch.no_grad():
+            if config.gate_rank_type == 1:
+                self.gate.weight.zero_()                
+            elif config.gate_rank_type == 2:
+                self.g1.zero_()
+                ortho_init(self.g2, 0.1)
+
             ratio_0_to_1 = layer_id / (n_layer - 1)  # 0 to 1
             ratio_1_to_almost0 = 1.0 - (layer_id / n_layer)  # 1 to ~0
-            ddd = torch.ones(1, 1, n_embd)
-            for i in range(n_embd):
-                ddd[0, 0, i] = i / n_embd
-
-            ddd = torch.zeros(1, 1, n_embd)
-            self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
-            self.time_maa_r = nn.Parameter(torch.zeros_like(ddd))
-            self.time_maa_k = nn.Parameter(torch.zeros_like(ddd))
-            self.time_maa_v = nn.Parameter(torch.zeros_like(ddd))
-            self.time_maa_w = nn.Parameter(torch.zeros_like(ddd))
-            self.time_maa_g = nn.Parameter(torch.zeros_like(ddd))
-
-            D_MIX_LORA = 32 if n_embd < 4096 else 64
-            self.time_maa_w2 = nn.Parameter(torch.zeros(5, D_MIX_LORA, n_embd).uniform_(-0.01, 0.01))
-            self.time_maa_w1 = nn.Parameter(torch.zeros(n_embd, D_MIX_LORA*self.time_maa_w2.size(0)))
+            
+            if self.config.use_tokenshift:
+                ddd = torch.ones(1, 1, n_embd)
+                for i in range(n_embd):
+                    ddd[0, 0, i] = i / n_embd
+
+                ddd = torch.zeros(1, 1, n_embd)
+                self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+                self.time_maa_r = nn.Parameter(torch.zeros_like(ddd))
+                self.time_maa_k = nn.Parameter(torch.zeros_like(ddd))
+                self.time_maa_v = nn.Parameter(torch.zeros_like(ddd))
+                self.time_maa_w = nn.Parameter(torch.zeros_like(ddd))
+                self.time_maa_g = nn.Parameter(torch.zeros_like(ddd))
+
+                lora_rank_tokenshift = config.lora_rank_tokenshift or (32 if n_embd < 4096 else 64)
+
+                self.time_maa_w2 = nn.Parameter(torch.zeros(5, lora_rank_tokenshift, n_embd).uniform_(-0.01, 0.01))
+                self.time_maa_w1 = nn.Parameter(torch.zeros(n_embd, lora_rank_tokenshift*self.time_maa_w2.size(0)))
+
+            lora_rank_decay = config.lora_rank_decay or (64 if n_embd < 4096 else 128)
 
             # RWKV-6
             decay_speed = torch.ones(dim_att)
             for n in range(dim_att):
                 decay_speed[n] = -6 + 5 * (n / (dim_att - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
             self.time_decay = nn.Parameter(decay_speed.reshape(1,1,dim_att))
-            D_DECAY_LORA = 64 if n_embd < 4096 else 128
-            self.time_decay_w1 = nn.Parameter(torch.zeros(n_embd, D_DECAY_LORA))
-            self.time_decay_w2 = nn.Parameter(torch.zeros(D_DECAY_LORA, dim_att).uniform_(-0.01, 0.01))
+            self.time_decay_w1 = nn.Parameter(torch.zeros(n_embd, lora_rank_decay))
+            self.time_decay_w2 = nn.Parameter(torch.zeros(lora_rank_decay, dim_att).uniform_(-0.01, 0.01))
 
     def forward(
         self,
@@ -291,11 +367,12 @@ class RWKV6Attention(nn.Module):
     ):
         output_shift_state = hidden_states[:, -1:].detach().clone()
 
-        bsz, q_len, hidden_dim = hidden_states.size()
-        H = self.num_heads
-
         x = hidden_states
 
+        B, T, C = hidden_states.shape
+        H = self.num_heads
+        N = self.head_dim
+
         if use_cache and past_key_values is not None and len(past_key_values) > self.layer_idx:
             input_kv_state, input_shift_state = past_key_values[self.layer_idx]
             xprev = torch.cat([input_shift_state, x[:, :-1]], dim=1)
@@ -303,83 +380,71 @@ class RWKV6Attention(nn.Module):
             input_kv_state = None
             xprev = F.pad(x, (0, 0, 1, -1))
 
-        dxprev = xprev - x
-
-        xxx = x + dxprev * self.time_maa_x
-        xxx = torch.tanh(xxx @ self.time_maa_w1).view(bsz*q_len, self.time_maa_w2.size(0), -1).transpose(0, 1)
-        xxx = torch.bmm(xxx, self.time_maa_w2).view(self.time_maa_w2.size(0), bsz, q_len, hidden_dim)
+        if self.config.use_tokenshift:
+            dxprev = xprev - x
 
-        mr, mk, mv, mw, mg = xxx.unbind(dim=0)
-        xr = x + dxprev * (self.time_maa_r + mr)
-        xk = x + dxprev * (self.time_maa_k + mk)
-        xv = x + dxprev * (self.time_maa_v + mv)
-        xw = x + dxprev * (self.time_maa_w + mw)
-        xg = x + dxprev * (self.time_maa_g + mg)
+            xxx = x + dxprev * self.time_maa_x
+            xxx = torch.tanh(xxx @ self.time_maa_w1).view(B*T, self.time_maa_w2.size(0), -1).transpose(0, 1)
+            xxx = torch.bmm(xxx, self.time_maa_w2).view(self.time_maa_w2.size(0), B, T, C)
 
-        query_states = self.q_proj(xr)
-        key_states = self.k_proj(xk)
-        value_states = self.v_proj(xv)
-        decay_states = (self.time_decay + torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2).to(query_states.dtype)
-        gate_states = F.sigmoid(self.gate(xg))
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        decay_states = decay_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+            mr, mk, mv, mw, mg = xxx.unbind(dim=0)
+            xr = x + dxprev * (self.time_maa_r + mr)
+            xk = x + dxprev * (self.time_maa_k + mk)
+            xv = x + dxprev * (self.time_maa_v + mv)
+            xw = x + dxprev * (self.time_maa_w + mw)
+            xg = x + dxprev * (self.time_maa_g + mg)
+        else:
+            xr = xk = xv = xw = xg = x
+
+        r = self.q_proj(xr)
+        k = self.k_proj(xk)
+        v = self.v_proj(xv)
+        w_lora_result = (self.time_decay + torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2).to(r.dtype)
+        if self.config.gate_rank_type == 1:
+            g = torch.sigmoid(self.gate(xg))
+        elif self.config.gate_rank_type == 2:
+            g = torch.sigmoid(xg @ self.g1) @ self.g2
+
+        if position_embeddings is not None:
+            r = r.view(B,T,-1,N)
+            k = k.view(B,T,-1,N)
+            cos, sin = position_embeddings
+            r, k = apply_rotary_pos_emb(r, k, cos, sin, unsqueeze_dim=2)
 
         # repeat k/v heads if n_kv_heads < n_heads
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        k = k.view(B, T, -1, 1, self.head_dim).expand(-1, -1, -1, self.num_key_value_groups, -1).reshape(B, T, -1)
+        v = v.view(B, T, -1, 1, self.head_dim).expand(-1, -1, -1, self.num_key_value_groups, -1).reshape(B, T, -1)
         dropout_rate = 0.0 if not self.training else self.attention_dropout
 
-        decay_states_log = -decay_states.float().exp()
-        decay_states_log = decay_states_log.clamp(-5) # FIXME - is this necessary?
-        key_states = (key_states * (1 - decay_states_log.exp())).to(key_states.dtype)
+        log_w = -w_lora_result.float().exp()
+        log_w = log_w.clamp(-5)
+        if self.config.balance_state:
+            k = (k * (1 - log_w.exp())).to(k.dtype)
 
+        # dealing with left-padding
         if attention_mask is not None:
-            if q_len > 1:
-                decay_states_log = decay_states_log - 100 * F.pad(1 - attention_mask, [1, -1]).view(bsz, 1, q_len, 1)
-
-        query_states = query_states.to(value_states.dtype)
-        key_states = key_states.to(value_states.dtype)
-
-        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
-        # therefore the input hidden states gets silently casted in float32. Hence, we need
-        # cast them back in float16 just to be sure everything works as expected.
-        input_dtype = query_states.dtype
-        if input_dtype == torch.float32:
-            if torch.is_autocast_enabled():
-                target_dtype = torch.get_autocast_gpu_dtype()
-            # Handle the case where the model is quantized
-            elif hasattr(self.config, "_pre_quantization_dtype"):
-                target_dtype = self.config._pre_quantization_dtype
-            else:
-                target_dtype = self.q_proj.weight.dtype
+            v = v * attention_mask[:, None, -v.shape[-2]:, None]
 
-            logger.warning_once(
-                f"The input hidden states seems to be silently casted in float32, this might be related to"
-                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
-                f" {target_dtype}."
-            )
-
-            query_states = query_states.to(target_dtype)
-            key_states = key_states.to(target_dtype)
-            value_states = value_states.to(target_dtype)
+        r = r.view(B,T,-1,N).to(v.dtype)
+        k = k.view(B,T,-1,N).to(v.dtype)
+        v = v.view(B,T,-1,N)
+        log_w = log_w.view(B,T,-1,N)
 
         attn_weights = torch.empty(0, device=x.device)
 
-        scale = query_states.shape[-1] ** -0.5
+        scale = r.shape[-1] ** -0.5
         output_final_state = not self.training and use_cache and past_key_values is not None
-        #attn_output, output_kv_state = ChunkGLAFunction.apply(query_states, key_states, value_states, decay_states_log.float(), scale, input_kv_state, output_final_state)
-        #attn_output, output_kv_state = chunk_gla(query_states, key_states, value_states, decay_states_log, scale, input_kv_state, output_final_state)
-        attn_output, output_kv_state = fused_recurrent_gla(query_states, key_states, value_states, decay_states_log, None, scale, input_kv_state, output_final_state)
+        attn_output, output_kv_state = fused_recurrent_gla(r, k, v, log_w, None, scale, input_kv_state, output_final_state)
 
         if output_final_state:
-            past_key_values.update(output_kv_state, output_shift_state, q_len, self.layer_idx)
+            past_key_values.update(output_kv_state, output_shift_state, T, self.layer_idx)
 
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(bsz, q_len, -1)
-        attn_output = self.o_proj(attn_output * gate_states)
+        attn_output = attn_output.view(B, T, -1)
+        if self.config.groupnorm_att:
+            attn_output = self.ln_x(attn_output.view(B * T, -1)).view(B, T, -1)
+        if self.config.gate_rank_type != 0:
+            attn_output = attn_output * g
+        attn_output = self.o_proj(attn_output)
 
         return attn_output, attn_weights
     
@@ -578,7 +643,7 @@ class RWKV6Qwen2Model(RWKV6Qwen2PreTrainedModel):
         )
         self._attn_implementation = config._attn_implementation
         self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        #self.rotary_emb = Qwen2RotaryEmbedding(config=config)
+        self.rotary_emb = Qwen2RotaryEmbedding(config=config)
 
         self.gradient_checkpointing = False
         # Initialize weights and apply final processing
@@ -640,13 +705,14 @@ class RWKV6Qwen2Model(RWKV6Qwen2PreTrainedModel):
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
-        # if cache_position is None:
-        #     past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        #     cache_position = torch.arange(
-        #         past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-        #     )
-        # if position_ids is None:
-        #     position_ids = cache_position.unsqueeze(0)
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
 
         # causal_mask = self._update_causal_mask(
         #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
@@ -657,7 +723,9 @@ class RWKV6Qwen2Model(RWKV6Qwen2PreTrainedModel):
         hidden_states = inputs_embeds
 
         # create position embeddings to be shared across the decoder layers
-        position_embeddings = None #self.rotary_emb(hidden_states, position_ids)
+        position_embeddings = None
+        if self.config.use_rope:
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         # decoder layers
         all_hidden_states = () if output_hidden_states else None