fla/modules/conv/cp/ops.py

import torch
import torch.distributed as dist

from fla.ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
from fla.ops.utils import prepare_chunk_indices


class CausalConv1dFunctionCP(torch.autograd.Function):
    """
    Context Parallel version of CausalConv1dFunction.

    Forward:
        1. Get tails from previous rank to construct initial_state
        2. Call causal_conv1d_fwd

    Backward:
        1. Call causal_conv1d_bwd to get dx
        2. Sync communication: add next rank's first W-1 token gradients to current rank's last W-1 tokens
    """

    @staticmethod
    def _prepare_initial_state_for_cp(
        x: torch.Tensor,
        weight: torch.Tensor,
        cu_seqlens: torch.Tensor | None,
        context: FLACPContext,
        group: dist.ProcessGroup | None,
    ) -> torch.Tensor | None:
        """Prepare initial_state for CP forward pass by communicating with previous rank.

        Args:
            x: Input tensor of shape [1, T, D]
            weight: Weight tensor of shape [D, W]
            cu_seqlens: Cumulative sequence lengths
            context: CP context
            group: Process group for communication

        Returns:
            initial_state: Initial state tensor of shape [N, D, W] or None
        """
        if group is None:
            return None

        W = weight.shape[-1]  # weight: [D, W]
        D = weight.shape[0]
        initial_state = None
        if not context.is_first_rank:
            # Non-first rank needs initial_state
            assert x.dim() == 3 and x.shape[0] == 1, f"CP requires [1, T, D], got {x.shape}"
            x_2d = x.squeeze(0)  # [T, D]
            tails = x_2d[-(W-1):].contiguous()  # [W-1, D]
            heads = conv_cp_send_recv_fwd(tails, group)  # [W-1, D]
            # Construct initial_state: [N, D, W]
            N = len(cu_seqlens) - 1
            initial_state = torch.zeros(N, D, W, device=x.device, dtype=x.dtype)
            valid_len = min(W - 1, context.pre_num_conv_tokens)
            if valid_len > 0:
                # heads[-valid_len:]: [valid_len, D] -> [D, valid_len]
                initial_state[0, :, -valid_len:] = heads[-valid_len:].T
        else:
            # First rank also needs to participate in communication (send tails)
            x_2d = x.squeeze(0)
            tails = x_2d[-(W-1):].contiguous()
            _ = conv_cp_send_recv_fwd(tails, group)  # Send but don't use

        return initial_state

    @staticmethod
    def _correct_dx_for_cp(
        dx: torch.Tensor,
        dh0: torch.Tensor | None,
        W: int,
        group: dist.ProcessGroup | None,
        is_first_rank: bool,
        pre_num_conv_tokens: int = 0,
    ) -> None:
        """Correct dx gradients for CP backward pass by communicating with next rank.

        Args:
            dx: Gradient tensor to be corrected, shape [1, T, D]
            dh0: Gradient w.r.t. initial_state, shape [N, D, W] or None
            W: Kernel size
            group: Process group for communication
            is_first_rank: Whether this is the first rank in the sequence's processing chain
            pre_num_conv_tokens: Number of tokens from the previous rank that
                belong to the first sequence on the current rank. Must match the
                value used in the forward pass to construct initial_state.
        """
        if group is None:
            return

        D = dx.shape[-1]
        # dh0: [N, D, W] or None
        # We only care about the first sequence's initial_state gradient
        if dh0 is not None:
            # Only keep gradients for positions that had real data from the
            # previous rank. The forward fills only the last valid_len positions
            # of initial_state; gradients for the remaining (zero-padded) positions
            # must not flow back, otherwise they leak into unrelated sequences.
            valid_len = min(W - 1, pre_num_conv_tokens)
            d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
            if valid_len > 0:
                d_initial_state[-valid_len:] = dh0[0, :, -valid_len:].T
        else:
            # dh0 is None only when this is the first rank (no initial_state needed)
            assert is_first_rank, "dh0 should not be None when is_first_rank=False"
            d_initial_state = torch.zeros(W-1, D, device=dx.device, dtype=dx.dtype)
        # Sync communication: send d_initial_state to previous rank, receive from next rank
        recv_d_init = conv_cp_send_recv_bwd(d_initial_state, group)  # [W-1, D]
        # Add to current rank's last W-1 tokens (these tokens are used as initial_state by next rank)
        dx[0, -(W-1):, :].add_(recv_d_init)

    @staticmethod
    def forward(
        ctx,
        x: torch.Tensor,
        weight: torch.Tensor,
        bias: torch.Tensor | None,
        activation: str | None,
        chunk_indices: torch.Tensor | None,
        cp_context: FLACPContext | None,
        chunk_size: int | None,
        backend: str = 'triton',
    ):
        # Import here to avoid circular dependency
        from fla.modules.conv.triton.ops import causal_conv1d_fwd

        if cp_context is None:
            raise ValueError("cp_context must be provided for CausalConv1dFunctionCP")
        cu_seqlens = cp_context.cu_seqlens
        cu_seqlens_cpu = cp_context.cu_seqlens_cpu
        group = cp_context.group

        # Get kernel_size
        W = weight.shape[-1]  # weight: [D, W]
        # Prepare initial_state for CP
        initial_state = CausalConv1dFunctionCP._prepare_initial_state_for_cp(
            x=x,
            weight=weight,
            cu_seqlens=cu_seqlens,
            context=cp_context,
            group=group,
        )

        ctx.save_for_backward(x, weight, bias, initial_state)
        ctx.activation = activation
        ctx.cu_seqlens = cu_seqlens
        ctx.cu_seqlens_cpu = cu_seqlens_cpu
        ctx.chunk_indices = chunk_indices
        ctx.chunk_size = chunk_size
        ctx.group = group
        ctx.W = W
        ctx.is_first_rank = cp_context.is_first_rank
        ctx.pre_num_conv_tokens = cp_context.pre_num_conv_tokens

        # Call original forward
        y, _ = causal_conv1d_fwd(
            x=x,
            weight=weight,
            bias=bias,
            residual=None,
            initial_state=initial_state,
            output_final_state=False,
            activation=activation,
            cu_seqlens=cu_seqlens,
            cu_seqlens_cpu=cu_seqlens_cpu,
            chunk_indices=chunk_indices,
            BT=chunk_size,
        )

        return y

    @staticmethod
    def backward(ctx, dy: torch.Tensor):
        # Import here to avoid circular dependency
        from fla.modules.conv.triton.ops import causal_conv1d_bwd

        x, weight, bias, initial_state = ctx.saved_tensors
        group = ctx.group
        W = ctx.W

        # Call original backward
        dx, dw, db, _, dh0 = causal_conv1d_bwd(
            x=x,
            dy=dy,
            dht=None,
            weight=weight,
            bias=bias,
            residual=None,
            initial_state=initial_state,
            activation=ctx.activation,
            cu_seqlens=ctx.cu_seqlens,
            cu_seqlens_cpu=ctx.cu_seqlens_cpu,
            chunk_indices=ctx.chunk_indices,
            BT=ctx.chunk_size,
        )

        # Correct dx gradients for CP
        CausalConv1dFunctionCP._correct_dx_for_cp(
            dx=dx,
            dh0=dh0,
            W=W,
            group=group,
            is_first_rank=ctx.is_first_rank,
            pre_num_conv_tokens=ctx.pre_num_conv_tokens,
        )

        return dx, dw, db, None, None, None, None, None


def causal_conv1d_cp(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    activation: str | None = None,
    chunk_indices: torch.Tensor | None = None,
    cp_context: FLACPContext | None = None,
    chunk_size: int | None = None,
    backend: str = 'triton',
):
    """
    Context Parallel version of causal_conv1d.

    Automatically handles communication in CP environment:
    - Forward: get initial_state from previous rank
    - Backward: correct dx gradients

    Args:
        x: Input tensor of shape [1, T, D]
        weight: Weight tensor of shape [D, W]
        bias: Bias tensor of shape [D] or None
        activation: Activation function name or None
        cu_seqlens: Cumulative sequence lengths
        cu_seqlens_cpu: Cumulative sequence lengths on CPU
        chunk_indices: Chunk indices for variable-length sequences
        cp_context: CP context (required for CP mode)
    """
    if cp_context is None:
        raise ValueError("cp_context must be provided for causal_conv1d_cp")

    assert cp_context.conv1d_kernel_size is not None, "conv1d_kernel_size must be provided for causal_conv1d_cp"
    assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for causal_conv1d_cp"
    assert backend in ['triton'], "backend must be 'triton'"
    chunk_size = chunk_size or 64
    if chunk_indices is None:
        chunk_indices = prepare_chunk_indices(cp_context.cu_seqlens, chunk_size, cu_seqlens_cpu=cp_context.cu_seqlens_cpu)

    return CausalConv1dFunctionCP.apply(
        x, weight, bias, activation,
        chunk_indices, cp_context, chunk_size, backend
    )