project_6/ixformer_sdk/train/functions/fused_rope.py

from typing import List, Tuple, Union

import ixformer._C as ops
import torch
from torch.autograd.function import Function

# adding by xuelu.peng 20240417
# from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_rope.py#L59
__all__ = ["fused_apply_rotary_pos_emb", "fused_apply_split_rotary_pos_emb", "fused_apply_rotary_pos_emb_cache"]


class FusedRoPEFunc(Function):
    """
    Fused RoPE function

    This implementation assumes the input tensor to be in `sbhd` format and the RoPE tensor to be
    of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive
    `.contiguous()` calls, thus it may not achieve the best memory access pattern.
    """

    @staticmethod
    def forward(
        ctx,
        t: torch.Tensor,
        freqs: torch.Tensor,
        transpose_output_memory: bool = False,
    ) -> torch.Tensor:
        # assert transpose_output_memory == False
        output = ops.train.fused_rope_forward(t, freqs, transpose_output_memory)
        ctx.save_for_backward(freqs)
        ctx.transpose_output_memory = transpose_output_memory

        return output

    @staticmethod
    def backward(
        ctx, grad_output: torch.Tensor
    ) -> Tuple[Union[torch.Tensor, None], ...]:

        (freqs,) = ctx.saved_tensors
        grad_input = ops.train.fused_rope_backward(
            grad_output, freqs, ctx.transpose_output_memory
        )
        return grad_input, None, None


class FusedFluxRoPEFunc(Function):
    """
    Fused FluxRoPE function

    This implementation assumes the input tensor to be in `bshd` format and the RoPE tensor to be
    of shape (s, d), and output shape is the same as input shape.
    """

    @staticmethod
    def forward(
        ctx,
        t: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        imp_mode : int = 1
    ) -> torch.Tensor:
        # assert transpose_output_memory == False
        output = ops.train.fused_rope_forward_cached(t, cos, sin, imp_mode)
        ctx.save_for_backward(cos, sin)
        return output

    @staticmethod
    def backward(
        ctx, grad_output: torch.Tensor
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        (cos, sin) = ctx.saved_tensors
        grad_input = ops.train.fused_rope_backward_cached(
            grad_output, cos, sin
        )
        return grad_input, None, None, None
        


def fused_apply_rotary_pos_emb(
    t: torch.Tensor,
    freqs: torch.Tensor,
    transpose_output_memory: bool = False,
) -> torch.Tensor:
    """Apply rotary positional embedding to input tensor T in `sbhd` format, where
    s: sequence length
    b: batch size
    h: head num
    d: dim of each head

    Args:
        t (Tensor): Input tensor T is of shape [s, b, h, d], dtype : torch.float32, torch.half
        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and
        `float` dtype
        transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b'
        dimension of the output's underlying memory format. This is very helpful when you want to
        get a contiguous tensor after calling `output.transpose(0, 1)`.

    Returns:
        Tensor: The input tensor after applying RoPE
    """
    return FusedRoPEFunc.apply(t, freqs, transpose_output_memory)


def fused_apply_rotary_pos_emb_cache(
    t: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    imp_mode: int = 1,
) -> torch.Tensor:
    """Apply rotary positional embedding to input tensor T in `bshd` format, where
    s: sequence length
    b: batch size
    h: head num
    d: dim of each head

    Args:
        t (Tensor): Input tensor T is of shape [b, s, h, d], dtype : torch.float32, torch.half, torch.bfloat16
        cos/sin (Tensor): Rotary Positional embedding tensor freq is of shape [s, d] and
        `float` dtype
        imp_mode (bool): Default to 1. 1 for flux/cogvideox/hunyuan-dit, img_mode=0 for Stable Audio. For now, only img_mode = 1 is supported.

    Returns:
        Tensor: The input tensor after applying RoPE
    """
    return FusedFluxRoPEFunc.apply(t, cos, sin, imp_mode)


class FusedSplitRoPEFunc(torch.autograd.Function):
    """
    Fused Split and RoPE function

    This implementation assumes the input tensor to be in `sbh3d` format and the RoPE tensor to be
    of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive
    `.contiguous()` calls, thus it may not achieve the best memory access pattern.

    input:    mix_q_k_v [s,b,hn_kv,h/hn_kv+2,d]
    output:   output_q, output_k, output_v [s,b,h,d]
    """
    @staticmethod
    def forward(
        ctx,
        mixed_q_k_v: torch.Tensor,
        freqs: torch.Tensor,
        transpose_output_memory: bool = False,
    ) -> torch.Tensor: 
        assert transpose_output_memory == False, "do not support transpose_output now"
        assert mixed_q_k_v.is_contiguous() == True, "mixed_q_k_v should be contiguous in FusedSplitRoPEFunc." 

        s, b, hn_kv, repplus2, d = mixed_q_k_v.size()
        num_key_value_groups = repplus2-2

        q, k, v = torch.split(mixed_q_k_v, (num_key_value_groups,1,1), dim=3)

        ctx.hn_kv = hn_kv
        output_q, output_k, output_v = torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d)

        ops.train.fused_split_rope_forward(
            q, k, v, freqs, output_q, output_k, output_v, transpose_output_memory, hn_kv, num_key_value_groups
        )

        ctx.save_for_backward(freqs)
        ctx.transpose_output_memory = transpose_output_memory

        return output_q, output_k, output_v

    @staticmethod
    def backward(
        ctx, grad_o_q: torch.Tensor, grad_o_k: torch.Tensor, grad_o_v: torch.Tensor
    ) -> Tuple[Union[torch.Tensor, None], ...]:
        # grad_o_q: [s,b,h,d]
        s,b,h,d = grad_o_q.size()

        hn_kv = ctx.hn_kv

        mixed_shape = (s, b, hn_kv,(h//hn_kv+2), d)

        if hn_kv == h:
            grad_mixed_q_k_v = torch.empty(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device,memory_format=torch.contiguous_format)  # torch.empty效率比torch.zeros高
        else:
            grad_mixed_q_k_v = torch.zeros(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device)  # 支持 gqa 的情况，kernel内需要进行累加，需要把qkv的梯度置零
        grad_q, grad_k, grad_v = torch.split(grad_mixed_q_k_v.view(s,b,hn_kv,-1,d), (h//hn_kv,1,1), dim=3)

        (freqs,) = ctx.saved_tensors
        ops.train.fused_split_rope_backward(
            grad_o_q, grad_o_k, grad_o_v, freqs, grad_q, grad_k, grad_v, ctx.transpose_output_memory
        )

        return grad_mixed_q_k_v, None, None

def fused_apply_split_rotary_pos_emb(
    mixed_q_k_v: torch.Tensor,
    freqs: torch.Tensor,
    transpose_output_memory: bool = False,
) -> torch.Tensor:
    """ Split mixed_q_k_v and apply rotary positional embedding to q and k in `sbhd` format, where
    s: sequence length
    b: batch size
    h: head num
    d: dim of each head
    hn_kv: num head of key and value

    Args:
        mixed_q_k_v (Tensor): Input tensor T is of shape  [s,b,hn_kv,h/hn_kv+2,d]
        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and
        `float` dtype
        transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b'
        dimension of the output's underlying memory format. This is very helpful when you want to
        get a contiguous tensor after calling `output.transpose(0, 1)`.

    Returns:
        Tensors: The input tensors after split and applying RoPE
    """
    return FusedSplitRoPEFunc.apply(mixed_q_k_v, freqs, transpose_output_memory)