[Iluvatar GPU] Optimize attention performance and fix moe load ckpt error (#3651)

fastdeploy/model_executor/layers/attention/iluvatar_attn_backend.py

@@ -16,13 +16,11 @@
 
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from math import sqrt
 from typing import TYPE_CHECKING, Optional
 
 import paddle
-from paddle.nn.functional.flash_attention import flash_attn_unpadded
 
 from fastdeploy.config import FDConfig
 from fastdeploy.model_executor.layers.attention.attention import Attention
@@ -30,7 +28,11 @@ from fastdeploy.model_executor.layers.attention.base_attention_backend import (
     AttentionBackend,
     AttentionMetadata,
 )
-from fastdeploy.model_executor.ops.iluvatar import paged_attention
+from fastdeploy.model_executor.ops.iluvatar import (
+    mixed_fused_paged_attention,
+    paged_attention,
+    prefill_fused_paged_attention,
+)
 
 if TYPE_CHECKING:
     from fastdeploy.model_executor.forward_meta import ForwardMeta
@@ -42,26 +44,7 @@ class IluvatarAttentionMetadata(AttentionMetadata):
     IluvatarAttentionMetadata
     """
 
-    # flash_attn metadata
-    cu_seqlens_q: Optional[paddle.Tensor] = None
-    cu_seqlens_k: Optional[paddle.Tensor] = None
-    fixed_seed_offset: Optional[paddle.Tensor] = None
-    attn_mask: Optional[paddle.Tensor] = None
-    attn_mask_start_row_indices: Optional[paddle.Tensor] = None
-    dropout: float = 0.0
-    causal: bool = True
-    return_softmax: bool = False
-    rng_name: str = ""
-
-    # paged_attn metadata
-    block_tables: Optional[paddle.Tensor] = None
-    seq_lens: Optional[paddle.Tensor] = None
-    num_kv_heads: int = 1
-    scale: float = 1.0
-    block_size: int = 1
-    max_context_len: int = 1
     alibi_slopes: Optional[paddle.Tensor] = None
-    # causal: bool = True
     window_left: int = -1
     window_right: int = -1
     softcap: float = 0.0
@@ -88,55 +71,44 @@ class IluvatarAttnBackend(AttentionBackend):
     def __init__(self, fd_config: FDConfig, kv_num_heads: int, num_heads: int, head_dim: int):
         super().__init__()
         self.attention_metadata = IluvatarAttentionMetadata()
-        self.attention_metadata.block_size = fd_config.parallel_config.block_size
-        assert (
-            fd_config.parallel_config.enc_dec_block_num == 0
-        ), f"Iluvatar does not support yet, {fd_config.parallel_config.enc_dec_block_num}"
-        assert self.attention_metadata.block_size == 16, "Iluvatar paged attn requires block_size must be 16."
-
-        self.attention_metadata.max_context_len = fd_config.parallel_config.max_model_len
-        self.attention_metadata.causal = getattr(fd_config.model_config, "causal", True)
+        self.block_size = fd_config.parallel_config.block_size
+        assert self.block_size == 16, "Iluvatar paged attn requires block_size must be 16."
+        self.max_context_len = fd_config.parallel_config.max_model_len
+        self.causal = getattr(fd_config.model_config, "causal", True)
         self.speculate_method = getattr(fd_config.parallel_config, "speculate_method", None)
         self.use_speculate = self.speculate_method is not None
-        self.attention_metadata.num_kv_heads = kv_num_heads
-        self.attention_metadata.dropout = fd_config.model_config.hidden_dropout_prob
+        self.num_kv_heads = kv_num_heads
         self.num_heads = num_heads
         self.total_num_heads = num_heads + 2 * kv_num_heads
         self.head_dim = head_dim
-        self.hidden_dim = num_heads * head_dim
-        self.total_hidden_dim = self.total_num_heads * head_dim
+        self.hidden_dim = fd_config.model_config.hidden_size
         # note: scale need to change if using MLA
-        self.attention_metadata.scale = 1.0 / sqrt(head_dim)
+        self.scale = 1.0 / sqrt(head_dim)
         self.num_layers = fd_config.model_config.num_hidden_layers
         self.dtype = paddle.get_default_dtype()
 
-        self.record_block_table_metadata = {}
-        self.enable_fused_attention = int(os.getenv("FD_ILUVATAR_ENABLE_FUSED_ATTN", 1))
-
     def init_attention_metadata(self, forward_meta: ForwardMeta):
         """Initialize attntion metadata hence all layers in the forward pass can reuse it."""
+        self.rope_cos = forward_meta.rotary_embs[0, 0, :, :, :]
+        self.rope_sin = forward_meta.rotary_embs[1, 0, :, :, :]
         self.prefill_info_dict = {}
         self.decode_info_dict = {}
-
-        prefill_non_zeros_ids = forward_meta.seq_lens_this_time > 1
-        decode_non_zeros_ids = forward_meta.seq_lens_this_time == 1
-        self.prefill_info_dict["batch_ids"] = paddle.where(prefill_non_zeros_ids)[0]
-        self.decode_info_dict["batch_ids"] = paddle.where(decode_non_zeros_ids)[0]
-
+        self.prefill_info_dict["batch_ids"] = paddle.where(forward_meta.seq_lens_encoder)[0]
+        self.decode_info_dict["batch_ids"] = paddle.where(forward_meta.seq_lens_decoder)[0]
         self.prefill_len = len(self.prefill_info_dict["batch_ids"])
         self.decode_len = len(self.decode_info_dict["batch_ids"])
         # only prefill
         if self.decode_len == 0:
             cu_seq_ids = list(range(self.prefill_len + 1))
             self.prefill_info_dict["cu_seqlens_q"] = forward_meta.cu_seqlens_q[cu_seq_ids]
+            self.mixed = False
         # only decode
         elif self.prefill_len == 0:
-            pass
+            self.mixed = False
         # both prefill and decode
         else:
-            prefill_num_tokens = paddle.sum(forward_meta.seq_lens_this_time[prefill_non_zeros_ids])
-            decode_num_tokens = paddle.sum(forward_meta.seq_lens_this_time[decode_non_zeros_ids])
-
+            self.mixed = True
+            self.prefill_num_tokens = paddle.sum(forward_meta.seq_lens_encoder).item()
             self.prefill_info_dict["cu_seqlens_q"] = paddle.zeros(
                 [self.prefill_len + 1], dtype=forward_meta.cu_seqlens_q.dtype
             )
@@ -145,36 +117,30 @@ class IluvatarAttnBackend(AttentionBackend):
             ]
             self.prefill_info_dict["cu_seqlens_q"] = paddle.cumsum(self.prefill_info_dict["cu_seqlens_q"])
 
-            self.prefill_qkv = paddle.zeros([prefill_num_tokens, self.total_hidden_dim], dtype=self.dtype)
-            self.decode_qkv = paddle.zeros([decode_num_tokens, self.total_hidden_dim], dtype=self.dtype)
-            self.merged_output = paddle.zeros(
-                [prefill_num_tokens + decode_num_tokens, self.num_heads, self.head_dim], dtype=self.dtype
+            self.tmp_buffer = paddle.zeros(
+                [self.prefill_num_tokens + self.decode_len, self.hidden_dim], dtype=self.dtype
             )
 
-            prefill_start, decode_start, start = 0, 0, 0
+            prefill_start, decode_start, start = 0, self.prefill_num_tokens, 0
             non_zeros_ids = forward_meta.seq_lens_this_time != 0
             non_zeros_seq_lens = forward_meta.seq_lens_this_time[non_zeros_ids]
             end = non_zeros_seq_lens[0]
             if end > 1:
                 last_stage = "prefill"
                 prefill_end = end
-                decode_end = 0
+                decode_end = decode_start
             else:
                 last_stage = "decode"
                 prefill_end = 0
-                decode_end = end
+                decode_end = decode_start + end
 
-            self.prefill_info_dict["id_group"] = []
-            self.prefill_info_dict["reverse_id_group"] = []
-            self.decode_info_dict["id_group"] = []
-            self.decode_info_dict["reverse_id_group"] = []
-            self.record_stages = []
+            self.id_group = []
+            self.reverse_id_group = []
             for seq_len in non_zeros_seq_lens[1:]:
                 if seq_len > 1:
                     if last_stage == "decode":
-                        self.record_stages.append((last_stage, len(self.decode_info_dict["id_group"])))
-                        self.decode_info_dict["id_group"].append((decode_start, decode_end))
-                        self.decode_info_dict["reverse_id_group"].append((start, end))
+                        self.id_group.append((decode_start, decode_end))
+                        self.reverse_id_group.append((start, end))
                         decode_start = decode_end
                         start = end
                         last_stage = "prefill"
@@ -182,9 +148,8 @@ class IluvatarAttnBackend(AttentionBackend):
                     end += seq_len
                 else:
                     if last_stage == "prefill":
-                        self.record_stages.append((last_stage, len(self.prefill_info_dict["id_group"])))
-                        self.prefill_info_dict["id_group"].append((prefill_start, prefill_end))
-                        self.prefill_info_dict["reverse_id_group"].append((start, end))
+                        self.id_group.append((prefill_start, prefill_end))
+                        self.reverse_id_group.append((start, end))
                         prefill_start = prefill_end
                         start = end
                         last_stage = "decode"
@@ -192,13 +157,11 @@ class IluvatarAttnBackend(AttentionBackend):
                     end += seq_len
 
             if prefill_start < prefill_end:
-                self.record_stages.append(("prefill", len(self.prefill_info_dict["id_group"])))
-                self.prefill_info_dict["id_group"].append((prefill_start, prefill_end))
-                self.prefill_info_dict["reverse_id_group"].append((start, end))
+                self.id_group.append((prefill_start, prefill_end))
+                self.reverse_id_group.append((start, end))
             if decode_start < decode_end:
-                self.record_stages.append(("decode", len(self.decode_info_dict["id_group"])))
-                self.decode_info_dict["id_group"].append((decode_start, decode_end))
-                self.decode_info_dict["reverse_id_group"].append((start, end))
+                self.id_group.append((decode_start, decode_end))
+                self.reverse_id_group.append((start, end))
 
     def get_attntion_meta(self):
         """get_attntion_meta"""
@@ -214,206 +177,20 @@ class IluvatarAttnBackend(AttentionBackend):
         """
         return (
             max_num_blocks,
-            self.attention_metadata.num_kv_heads,
-            self.attention_metadata.block_size,
+            self.num_kv_heads,
+            self.block_size,
             self.head_dim,
         )
 
-    def prefill_update_kv_cache(
-        self, k, v, k_cache_id: int, v_cache_id: int, layer_id: int, forward_meta: ForwardMeta, prefill_batch_ids: list
-    ):
-        # [num_tokens, num_kv_heads, head_dim] -> [num_kv_heads, num_tokens, head_dim]
-        trans_k = k.transpose([1, 0, 2]).contiguous()
-        trans_v = v.transpose([1, 0, 2]).contiguous()
-        tensor_start = 0
-        for batch_idx in prefill_batch_ids:
-            seq_len = forward_meta.seq_lens_this_time[batch_idx]
+    def transpose(self, hidden_states):
+        for ids, reverse_ids in zip(self.id_group, self.reverse_id_group):
+            self.tmp_buffer[ids[0] : ids[1], :] = hidden_states[reverse_ids[0] : reverse_ids[1], :]
+        return self.tmp_buffer
 
-            tensor_end = tensor_start + seq_len
-            slice_trans_k = trans_k[:, tensor_start:tensor_end, :]
-            slice_trans_v = trans_v[:, tensor_start:tensor_end, :]
-
-            cur_block_tables = forward_meta.block_tables[batch_idx]
-            cur_used_block_tables = cur_block_tables[cur_block_tables != -1]
-
-            cache_start = 0
-            cur_used_num_blocks = cur_used_block_tables.shape[0]
-            for i, block_id in enumerate(cur_used_block_tables):
-                # last block: seq_len - cache_start <= block_size
-                if i == cur_used_num_blocks - 1:
-                    cache_end = seq_len - cache_start
-                    assert cache_end <= self.attention_metadata.block_size
-                    paddle.assign(
-                        slice_trans_k[:, cache_start:seq_len, :],
-                        output=forward_meta.caches[k_cache_id][block_id, :, 0:cache_end, :],
-                    )
-                    paddle.assign(
-                        slice_trans_v[:, cache_start:seq_len, :],
-                        output=forward_meta.caches[v_cache_id][block_id, :, 0:cache_end, :],
-                    )
-                    if layer_id == self.num_layers - 1:
-                        self.record_block_table_metadata[batch_idx] = {
-                            "block_id": block_id.item(),
-                            "cache_end": cache_end.item(),
-                        }
-                # non last block: seq_lens_this_time > block_size
-                else:
-                    assert seq_len > self.attention_metadata.block_size
-                    cache_end = cache_start + self.attention_metadata.block_size
-                    paddle.assign(
-                        slice_trans_k[:, cache_start:cache_end, :], output=forward_meta.caches[k_cache_id][block_id]
-                    )
-                    paddle.assign(
-                        slice_trans_v[:, cache_start:cache_end, :], output=forward_meta.caches[v_cache_id][block_id]
-                    )
-                    cache_start += self.attention_metadata.block_size
-
-            tensor_start = tensor_end
-
-    def get_splited_qkv(
-        self, qkv: paddle.Tensor, forward_meta: ForwardMeta, cu_seqlens_q: paddle.Tensor, batch_ids=None
-    ):
-        q_end = self.hidden_dim
-        k_end = q_end + self.attention_metadata.num_kv_heads * self.head_dim
-        v_end = k_end + self.attention_metadata.num_kv_heads * self.head_dim
-        assert v_end == qkv.shape[-1], f"Shape mismatch: {v_end} vs {qkv.shape[-1]}"
-        assert qkv.shape[0] == cu_seqlens_q[-1], f"Shape mismatch: {qkv.shape[0]} vs {cu_seqlens_q[-1]}"
-
-        if batch_ids is None:
-            batch_ids = list(range(forward_meta.seq_lens_this_time.shape[0]))
-
-        q = qkv[..., 0:q_end]
-        k = qkv[..., q_end:k_end]
-        v = qkv[..., k_end:v_end]
-        q = q.view([-1, self.num_heads, self.head_dim])
-        k = k.view([-1, self.attention_metadata.num_kv_heads, self.head_dim])
-        v = v.view([-1, self.attention_metadata.num_kv_heads, self.head_dim])
-
-        for idx in range(len(cu_seqlens_q) - 1):
-            batch_idx = batch_ids[idx]
-            seq_len_i = forward_meta.seq_lens_this_time[batch_idx]
-            if seq_len_i == 0:
-                continue
-            cached_kv_len = forward_meta.seq_lens_decoder[batch_idx][0]
-            cu_seq_start_q = cu_seqlens_q[idx]
-            cu_seq_end_q = cu_seqlens_q[idx + 1]
-            # forward_meta.rotary_embs is [2, 1, S, 1, D]
-            if forward_meta.rotary_embs is not None:
-                cos = forward_meta.rotary_embs[0, 0, cached_kv_len : cached_kv_len + seq_len_i, :, :]
-                sin = forward_meta.rotary_embs[1, 0, cached_kv_len : cached_kv_len + seq_len_i, :, :]
-                q[cu_seq_start_q:cu_seq_end_q] = apply_rope(q[cu_seq_start_q:cu_seq_end_q], cos, sin)
-                k[cu_seq_start_q:cu_seq_end_q] = apply_rope(k[cu_seq_start_q:cu_seq_end_q], cos, sin)
-
-        return q, k, v
-
-    def split_pd_qkv(self, qkv):
-
-        for ids, reverse_ids in zip(self.prefill_info_dict["id_group"], self.prefill_info_dict["reverse_id_group"]):
-            self.prefill_qkv[ids[0] : ids[1], :] = qkv[reverse_ids[0] : reverse_ids[1], :]
-
-        for ids, reverse_ids in zip(self.decode_info_dict["id_group"], self.decode_info_dict["reverse_id_group"]):
-            self.decode_qkv[ids[0] : ids[1], :] = qkv[reverse_ids[0] : reverse_ids[1], :]
-
-        return self.prefill_qkv, self.decode_qkv
-
-    def merge_pd_output(self, prefill_out, decode_out):
-        for stage, idx in self.record_stages:
-            if stage == "prefill":
-                ids = self.prefill_info_dict["id_group"][idx]
-                reverse_ids = self.prefill_info_dict["reverse_id_group"][idx]
-                self.merged_output[reverse_ids[0] : reverse_ids[1], :, :] = prefill_out[ids[0] : ids[1], :, :]
-            else:
-                ids = self.decode_info_dict["id_group"][idx]
-                reverse_ids = self.decode_info_dict["reverse_id_group"][idx]
-                self.merged_output[reverse_ids[0] : reverse_ids[1], :, :] = decode_out[ids[0] : ids[1], :, :]
-        return self.merged_output
-
-    def forward_prefill(self, prefill_qkv, layer_id, k_cache_id, v_cache_id, forward_meta: ForwardMeta):
-        prefill_q, prefill_k, prefill_v = self.get_splited_qkv(
-            prefill_qkv,
-            forward_meta,
-            self.prefill_info_dict["cu_seqlens_q"],
-            batch_ids=self.prefill_info_dict["batch_ids"],
-        )
-
-        prefill_out = flash_attn_unpadded(
-            prefill_q,
-            prefill_k,
-            prefill_v,
-            cu_seqlens_q=self.prefill_info_dict["cu_seqlens_q"],
-            cu_seqlens_k=self.prefill_info_dict["cu_seqlens_q"],
-            max_seqlen_q=self.attention_metadata.max_context_len,
-            max_seqlen_k=self.attention_metadata.max_context_len,
-            scale=self.attention_metadata.scale,
-            dropout=self.attention_metadata.dropout,
-            causal=self.attention_metadata.causal,
-            return_softmax=self.attention_metadata.return_softmax,
-        )[0]
-        self.prefill_update_kv_cache(
-            prefill_k, prefill_v, k_cache_id, v_cache_id, layer_id, forward_meta, self.prefill_info_dict["batch_ids"]
-        )
-
-        return prefill_out
-
-    def forward_decode(self, decode_qkv, k_cache_id, v_cache_id, forward_meta: ForwardMeta):
-        k_cache = forward_meta.caches[k_cache_id]
-        v_cache = forward_meta.caches[v_cache_id]
-        if self.enable_fused_attention:
-            rope_cos = forward_meta.rotary_embs[0, 0, :, :, :]
-            rope_sin = forward_meta.rotary_embs[1, 0, :, :, :]
-            decode_out = paged_attention(
-                decode_qkv.view([-1, self.total_num_heads, self.head_dim]),
-                k_cache,
-                v_cache,
-                block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
-                seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
-                num_kv_heads=self.attention_metadata.num_kv_heads,
-                scale=self.attention_metadata.scale,
-                block_size=self.attention_metadata.block_size,
-                max_context_len=self.attention_metadata.max_context_len,
-                alibi_slopes=self.attention_metadata.alibi_slopes,
-                causal=self.attention_metadata.causal,
-                window_left=self.attention_metadata.window_left,
-                window_right=self.attention_metadata.window_right,
-                softcap=self.attention_metadata.softcap,
-                use_cuda_graph=self.attention_metadata.use_cuda_graph,
-                use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
-                merged_qkv=True,
-                k=decode_qkv,
-                v=decode_qkv,
-                rope_sin=rope_sin,
-                rope_cos=rope_cos,
-            )
-        else:
-            decode_q, decode_k, decode_v = self.get_splited_qkv(
-                decode_qkv,
-                forward_meta,
-                self.decode_info_dict["cu_seqlens_q"],
-                batch_ids=self.decode_info_dict["batch_ids"],
-            )
-
-            decode_out = paged_attention(
-                decode_q,
-                k_cache,
-                v_cache,
-                block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
-                seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
-                num_kv_heads=self.attention_metadata.num_kv_heads,
-                scale=self.attention_metadata.scale,
-                block_size=self.attention_metadata.block_size,
-                max_context_len=self.attention_metadata.max_context_len,
-                alibi_slopes=self.attention_metadata.alibi_slopes,
-                causal=self.attention_metadata.causal,
-                window_left=self.attention_metadata.window_left,
-                window_right=self.attention_metadata.window_right,
-                softcap=self.attention_metadata.softcap,
-                use_cuda_graph=self.attention_metadata.use_cuda_graph,
-                use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
-                k=decode_k,
-                v=decode_v,
-            )
-
-        return decode_out
+    def reverse_transpose(self, hidden_states):
+        for ids, reverse_ids in zip(self.id_group, self.reverse_id_group):
+            self.tmp_buffer[reverse_ids[0] : reverse_ids[1], :] = hidden_states[ids[0] : ids[1], :]
+        return self.tmp_buffer
 
     def forward_mixed(
         self,
@@ -429,23 +206,84 @@ class IluvatarAttnBackend(AttentionBackend):
         """
         forward_mixed
         """
-        assert not self.use_speculate, "IluvatarAttnBackend cannot support speculate now"
         layer_id = layer.layer_id
         k_cache_id = layer_id * 2
         v_cache_id = k_cache_id + 1
-        q_dim = qkv.dim()
-        assert q_dim == 2
-
+        k_cache = forward_meta.caches[k_cache_id]
+        v_cache = forward_meta.caches[v_cache_id]
         if self.decode_len == 0:
-            output = self.forward_prefill(qkv, layer_id, k_cache_id, v_cache_id, forward_meta)
-
+            output = prefill_fused_paged_attention(
+                qkv,
+                k_cache,
+                v_cache,
+                block_tables=forward_meta.block_tables[self.prefill_info_dict["batch_ids"], :],
+                cu_seqlens_qkv=self.prefill_info_dict["cu_seqlens_q"],
+                num_heads=self.num_heads,
+                head_dim=self.head_dim,
+                num_kv_heads=self.num_kv_heads,
+                block_size=self.block_size,
+                max_seq_len=self.max_context_len,
+                scale=self.scale,
+                causal=self.causal,
+                q_rope=True,
+                k_rope=True,
+                v_rope=False,
+                rope_sin=self.rope_sin,
+                rope_cos=self.rope_cos,
+            )
         elif self.prefill_len == 0:
-            output = self.forward_decode(qkv, k_cache_id, v_cache_id, forward_meta)
+            output = paged_attention(
+                qkv,
+                k_cache,
+                v_cache,
+                block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
+                seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
+                num_heads=self.num_heads,
+                head_dim=self.head_dim,
+                num_kv_heads=self.num_kv_heads,
+                scale=self.scale,
+                block_size=self.block_size,
+                max_context_len=self.max_context_len,
+                alibi_slopes=self.attention_metadata.alibi_slopes,
+                causal=self.causal,
+                window_left=self.attention_metadata.window_left,
+                window_right=self.attention_metadata.window_right,
+                softcap=self.attention_metadata.softcap,
+                use_cuda_graph=self.attention_metadata.use_cuda_graph,
+                use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
+                merged_qkv=True,
+                k=qkv,
+                v=qkv,
+                rope_sin=self.rope_sin,
+                rope_cos=self.rope_cos,
+            )
         else:
-            prefill_qkv, decode_qkv = self.split_pd_qkv(qkv)
-            prefill_output = self.forward_prefill(prefill_qkv, layer_id, k_cache_id, v_cache_id, forward_meta)
-            decode_output = self.forward_decode(decode_qkv, k_cache_id, v_cache_id, forward_meta)
-            output = self.merge_pd_output(prefill_output, decode_output)
+            output = mixed_fused_paged_attention(
+                qkv,
+                k_cache,
+                v_cache,
+                prefill_block_tables=forward_meta.block_tables[self.prefill_info_dict["batch_ids"], :],
+                decode_block_tables=forward_meta.block_tables[self.decode_info_dict["batch_ids"], :],
+                cu_seqlens_qkv=self.prefill_info_dict["cu_seqlens_q"],
+                seq_lens=forward_meta.seq_lens_decoder[self.decode_info_dict["batch_ids"], 0] + 1,
+                prefill_num_tokens=self.prefill_num_tokens,
+                num_heads=self.num_heads,
+                head_dim=self.head_dim,
+                num_kv_heads=self.num_kv_heads,
+                block_size=self.block_size,
+                max_seq_len=self.max_context_len,
+                scale=self.scale,
+                causal=self.causal,
+                q_rope=True,
+                k_rope=True,
+                v_rope=False,
+                window_left=self.attention_metadata.window_left,
+                window_right=self.attention_metadata.window_right,
+                softcap=self.attention_metadata.softcap,
+                use_cuda_graph=self.attention_metadata.use_cuda_graph,
+                use_sqrt_alibi=self.attention_metadata.use_sqrt_alibi,
+                rope_sin=self.rope_sin,
+                rope_cos=self.rope_cos,
+            )
 
-        output = output.view([-1, self.num_heads * self.head_dim])
         return output