vllm.attention.backends.differential_flash_attn

class DifferentialFlashAttentionBackend(AttentionBackend):
    accept_output_buffer = False

    @staticmethod
    def get_supported_head_sizes() -> List[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> Tuple[int, ...]:
        if block_size % 16 != 0:
            raise ValueError("Block size must be a multiple of 16.")
        assert num_kv_heads % 2 == 0, "num_kv_heads must be divisible by 2"
        return (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size)

    @staticmethod
    def get_name() -> str:
        return "DIFFERENTIAL_FLASH_ATTN"

    @staticmethod
    def get_impl_cls() -> Type["DifferentialFlashAttentionImpl"]:
        return DifferentialFlashAttentionImpl

    @staticmethod
    def get_metadata_cls() -> Type["DifferentialFlashAttentionMetadata"]:
        return DifferentialFlashAttentionMetadata

    @staticmethod
    def get_builder_cls() -> Type["DifferentialFlashAttentionMetadataBuilder"]:
        return DifferentialFlashAttentionMetadataBuilder

    @staticmethod
    def get_state_cls() -> Type["CommonAttentionState"]:
        return CommonAttentionState

    @staticmethod
    def swap_blocks(
        src_kv_cache: torch.Tensor,
        dst_kv_cache: torch.Tensor,
        src_to_dst: torch.Tensor,
    ) -> None:
        src_key_cache = src_kv_cache[0]
        dst_key_cache = dst_kv_cache[0]
        ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
        src_value_cache = src_kv_cache[1]
        dst_value_cache = dst_kv_cache[1]
        ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

    @staticmethod
    def copy_blocks(
        kv_caches: List[torch.Tensor],
        src_to_dists: torch.Tensor,
    ) -> None:
        key_caches = [kv_cache[0] for kv_cache in kv_caches]
        value_caches = [kv_cache[1] for kv_cache in kv_caches]

        ops.copy_blocks(key_caches, value_caches, src_to_dists)

class DifferentialFlashAttentionImpl(AttentionImpl):
    """
    If the input tensors contain prompt tokens, the layout is as follows:
    |<--------------- num_prefill_tokens ----------------->|	
    |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|

    Otherwise, the layout is as follows:	
    |<----------------- num_decode_tokens ------------------>|	
    |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|

    Generation tokens can contain padding when cuda-graph is used.
    Currently, prompt tokens don't contain any padding.

    The prompts might have different lengths, while the generation tokens
    always have length 1.

    If chunked prefill is enabled, prefill tokens and decode tokens can be
    batched together in a flattened 1D query.

    |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
    |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

    Currently, cuda graph is disabled for chunked prefill, meaning there's no
    padding between prefill and decode tokens.
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: Optional[List[float]],
        sliding_window: Optional[int],
        kv_cache_dtype: str,
        logits_soft_cap: Optional[float] = None,
        attn_type: str = AttentionType.DECODER,
        kv_sharing_target_layer_name: Optional[str] = None,
        use_irope: bool = False,
        differential_flash_attention_config: Optional[Dict[str, Any]] = None,
    ) -> None:
        if differential_flash_attention_config is None:
            differential_flash_attention_config = {}
        self.differential_flash_attention_config = \
            differential_flash_attention_config
        self.used_shared_kv_cache = kv_sharing_target_layer_name is not None
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
        if use_irope:
            logger.warning(
                "Using irope in V0 is not supported yet, it will fall back "
                "to global attention for long context.")
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        if alibi_slopes is not None:
            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
        self.alibi_slopes = alibi_slopes
        self.sliding_window = ((sliding_window - 1,
                                0) if sliding_window is not None else (-1, -1))
        self.kv_cache_dtype = kv_cache_dtype
        self.vllm_flash_attn_version = get_flash_attn_version(
            requires_alibi=self.alibi_slopes is not None)
        if is_quantized_kv_cache(self.kv_cache_dtype) and (
                not self.kv_cache_dtype.startswith("fp8")
                or not flash_attn_supports_fp8()):
            raise NotImplementedError(
                f"FlashAttention does not support {self.kv_cache_dtype} "
                "kv-cache on this device "
                f"(FA supports fp8 = {flash_attn_supports_fp8()}).")
        if logits_soft_cap is None:
            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
            logits_soft_cap = 0
        self.logits_soft_cap = logits_soft_cap

        assert self.num_heads % self.num_kv_heads == 0
        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

        support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()
        if head_size not in support_head_sizes:
            raise ValueError(
                f"Head size {head_size} is not supported by FlashAttention. "
                f"Supported head sizes are: {support_head_sizes}.")
        self.attn_type = attn_type

        self.lambda_full = None
        self.subln = self.differential_flash_attention_config["subln"]

    def split_heads(self, x):
        # split by num_heads, the stripe pattern is friendly to tensor parallel.
        x = rearrange(x, "... (H two) D -> ... H two D", two=2)
        x1 = x[..., 0, :]
        x2 = x[..., 1, :]
        return x1.contiguous(), x2.contiguous()

    def split_kv_cache(self, x):
        # split by num_heads, the stripe pattern is friendly to tensor parallel.
        if x.numel() == 0:
            return torch.empty(0), torch.empty(0)

        x1, x2 = x[0], x[1]
        return x1, x2

    def populate_kv_cache(self, layer: AttentionLayer, key: torch.Tensor,
                          value: torch.Tensor, kv_cache: torch.Tensor,
                          attn_metadata: DifferentialFlashAttentionMetadata):
        if kv_cache.numel() > 0 and key is not None and value is not None:
            updated_slot_mapping = attn_metadata.slot_mapping
            torch.ops._C_cache_ops.reshape_and_cache_flash(
                key,
                value,
                kv_cache[0],
                kv_cache[1],
                updated_slot_mapping.flatten(),
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
            )

    def forward_generate_kv_cache(
            self, query: torch.Tensor, key: Optional[torch.Tensor],
            value: Optional[torch.Tensor], k_cache: torch.Tensor,
            v_cache: torch.Tensor,
            attn_metadata: DifferentialFlashAttentionMetadata) -> torch.Tensor:

        head_size = self.head_size
        num_heads = self.num_heads // 2
        num_kv_heads = self.num_kv_heads // 2

        query = query.view(-1, num_heads, head_size)
        if key is not None:
            assert value is not None
            key = key.view(-1, num_kv_heads, head_size)
            value = value.view(-1, num_kv_heads, head_size)
        else:
            assert value is None

        num_prefill_tokens = attn_metadata.num_prefill_tokens
        num_decode_tokens = attn_metadata.num_decode_tokens
        assert key.shape[
            0] == num_prefill_tokens + num_decode_tokens, "key shape mismatch"
        assert value.shape[
            0] == num_prefill_tokens + num_decode_tokens, "value shape mismatch"

        output = torch.empty_like(query)
        # Query for decode. KV is not needed because it is already cached.
        decode_query = query[num_prefill_tokens:]
        # QKV for prefill.
        query = query[:num_prefill_tokens]
        if key is not None and value is not None:
            key = key[:num_prefill_tokens]
            value = value[:num_prefill_tokens]

        assert query.shape[0] == num_prefill_tokens, "query shape mismatch"
        assert decode_query.shape[
            0] == num_decode_tokens, "decode query shape mismatch"

        if prefill_meta := attn_metadata.prefill_metadata:
            # Prompt run.
            if k_cache.numel() == 0 \
                or prefill_meta.block_tables is None \
                or prefill_meta.block_tables.numel() == 0:
                # normal attention
                prefill_output = flash_attn_varlen_func(
                    q=query,
                    k=key,
                    v=value,
                    cu_seqlens_q=prefill_meta.seq_start_loc,
                    cu_seqlens_k=prefill_meta.seq_start_loc,
                    max_seqlen_q=prefill_meta.max_prefill_seq_len,
                    max_seqlen_k=prefill_meta.max_prefill_seq_len,
                    softmax_scale=self.scale,
                    causal=True,
                    window_size=self.sliding_window,
                    alibi_slopes=self.alibi_slopes,
                    softcap=self.logits_soft_cap,
                )
                assert prefill_output.shape == output[:
                                                      num_prefill_tokens].shape
                output[:num_prefill_tokens] = prefill_output
            else:
                raise Exception("prefix caching not supported")

        if decode_meta := attn_metadata.decode_metadata:
            block_tables_arg = decode_meta.block_tables
            try:
                output[num_prefill_tokens:] = flash_attn_with_kvcache(
                    q=decode_query.unsqueeze(1),
                    k_cache=k_cache,
                    v_cache=v_cache,
                    block_table=block_tables_arg,
                    cache_seqlens=decode_meta.seq_lens_tensor,
                    softmax_scale=self.scale,
                    causal=True,
                    window_size=self.sliding_window,
                    alibi_slopes=self.alibi_slopes,
                    softcap=self.logits_soft_cap,
                ).squeeze(1)
            except Exception as e:
                logger.error("Error in PagedAttention.forward_decode: %s",
                             str(e))
                raise e

        # Reshape the output tensor.
        return output.view(-1, num_heads, head_size)

    def forward_with_kv_cache_only(
        self,
        query: torch.Tensor,
        k_cache: torch.Tensor,
        v_cache: torch.Tensor,
        attn_metadata: DifferentialFlashAttentionMetadata,
    ):
        if not attn_metadata.decode_metadata:
            block_tables_arg = attn_metadata.cross_layer_shared_block_tables
        else:
            block_tables_arg = attn_metadata.block_tables

        output = flash_attn_with_kvcache(
            q=query.unsqueeze(1),
            k_cache=k_cache,
            v_cache=v_cache,
            block_table=block_tables_arg,
            cache_seqlens=attn_metadata.seq_lens_tensor,
            softmax_scale=self.scale,
            causal=True,
            window_size=self.sliding_window,
            alibi_slopes=self.alibi_slopes,
            softcap=self.logits_soft_cap,
        ).squeeze(1)
        return output

    def forward(
        self,
        layer: AttentionLayer,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: DifferentialFlashAttentionMetadata,
        output: Optional[torch.Tensor] = None,
        output_scale: Optional[torch.Tensor] = None,
        output_block_scale: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass with FlashAttention.

        Args:
            query: shape = [num_tokens, num_heads, head_size]
            key: shape = [num_tokens, num_kv_heads, head_size]
            value: shape = [num_tokens, num_kv_heads, head_size]
            output: shape = [num_tokens, num_heads, head_size]
            kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
                NOTE: kv_cache will be an empty tensor with shape [0]
                for profiling run.
            attn_metadata: Metadata for attention.
        NOTE: It in-place updates the output tensor.
        NOTE: FP8 quantization, flash-attn expect the size of
              {q,k,v}_descale to be (num_sequences, num_kv_heads).
              We use torch's .expand() to avoid duplicating values
        """
        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError(
                "fused output quantization is not yet supported"
                " for DifferentialFlashAttentionImpl")

        if self.lambda_full is None:
            self.lambda_init = self.differential_flash_attention_config[
                "lambda_init"]
            lambda_q1 = self.differential_flash_attention_config["lambda_q1"]
            lambda_k1 = self.differential_flash_attention_config["lambda_k1"]
            lambda_q2 = self.differential_flash_attention_config["lambda_q2"]
            lambda_k2 = self.differential_flash_attention_config["lambda_k2"]
            lambda_1 = torch.exp(
                torch.sum(lambda_q1 * lambda_k1, dim=-1).float()).type_as(q)
            lambda_2 = torch.exp(
                torch.sum(lambda_q2 * lambda_k2, dim=-1).float()).type_as(q)
            self.lambda_full = lambda_1 - lambda_2 + self.lambda_init

        if not self.used_shared_kv_cache:  # need to generate kv-cache
            q = q.view(-1, self.num_heads, self.head_size)
            k = k.view(-1, self.num_kv_heads, self.head_size)
            v = v.view(-1, self.num_kv_heads, self.head_size)

            q1, q2 = self.split_heads(q)
            k1, k2 = self.split_heads(k)
            v1, v2 = self.split_heads(v)

            # kv_cache shape is (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size) # noqa: E501
            # Split by half along the first dimension.
            kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)
            assert kv_cache1.is_contiguous(), "kv_cache1 is not contiguous"
            assert kv_cache2.is_contiguous(), "kv_cache2 is not contiguous"

            if kv_cache1.numel() != 0:
                self.populate_kv_cache(layer, k1, v1, kv_cache1, attn_metadata)
                self.populate_kv_cache(layer, k2, v2, kv_cache2, attn_metadata)

                key_cache1, value_cache1 = self.split_kv_cache(kv_cache1)
                key_cache2, value_cache2 = self.split_kv_cache(kv_cache2)
            else:
                key_cache1, value_cache1 = torch.empty(0), torch.empty(0)
                key_cache2, value_cache2 = torch.empty(0), torch.empty(0)
            attn11 = self.forward_generate_kv_cache(q1, k1, v1, key_cache1,
                                                    value_cache1,
                                                    attn_metadata)
            attn12 = self.forward_generate_kv_cache(q1, k1, v2, key_cache1,
                                                    value_cache2,
                                                    attn_metadata)
            attn11 = attn11.view(q1.shape)
            attn12 = attn12.view(q1.shape)
            attn1 = torch.cat([attn11, attn12], dim=-1)

            attn21 = self.forward_generate_kv_cache(q2, k2, v1, key_cache2,
                                                    value_cache1,
                                                    attn_metadata)
            attn22 = self.forward_generate_kv_cache(q2, k2, v2, key_cache2,
                                                    value_cache2,
                                                    attn_metadata)
            attn21 = attn21.view(q2.shape)
            attn22 = attn22.view(q2.shape)
            attn2 = torch.cat([attn21, attn22], dim=-1)

            attn = attn1 - self.lambda_full * attn2
            # attn shape (-1, self.num_heads // 2, 2 * self.head_dim)
            attn = self.subln(attn)
            attn = attn * (1 - self.lambda_init)
            # reshape back to 2 * num_head
            attn_output = rearrange(attn,
                                    "... H (two D) -> ... (H two) D",
                                    two=2)

        else:  # reuse the kv cache, full attention
            q = q.view(-1, self.num_heads, self.head_size)
            q1, q2 = self.split_heads(q)
            # kv_cache shape is (2, num_blocks, block_size, num_kv_heads, head_size) # noqa: E501
            kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)
            key_cache1, value_cache1 = kv_cache1[0], kv_cache1[1]
            key_cache2, value_cache2 = kv_cache2[0], kv_cache2[1]

            attn11 = self.forward_with_kv_cache_only(q1, key_cache1,
                                                     value_cache1,
                                                     attn_metadata)
            attn12 = self.forward_with_kv_cache_only(q1, key_cache1,
                                                     value_cache2,
                                                     attn_metadata)
            attn11 = attn11.view(q1.shape)
            attn12 = attn12.view(q1.shape)
            attn1 = torch.cat([attn11, attn12], dim=-1)

            attn21 = self.forward_with_kv_cache_only(q2, key_cache2,
                                                     value_cache1,
                                                     attn_metadata)
            attn22 = self.forward_with_kv_cache_only(q2, key_cache2,
                                                     value_cache2,
                                                     attn_metadata)
            attn21 = attn21.view(q2.shape)
            attn22 = attn22.view(q2.shape)
            attn2 = torch.cat([attn21, attn22], dim=-1)

            attn = attn1 - self.lambda_full * attn2
            attn = self.subln(attn)
            attn = attn * (1 - self.lambda_init)
            # reshape back to 2 * num_head
            attn_output = rearrange(attn,
                                    "... H (two D) -> ... (H two) D",
                                    two=2)
        attn_output = attn_output.view(-1, self.num_heads * self.head_size)
        return attn_output