Hugging Face's logo

Kernels:
kernels-community
/
fla
Trusted publisher

Kernel card Files
xet
 Community
fla / build /torch-cuda /modules /conv /cp /ops.py
kernels-bot's picture
kernels-bot
Uploaded using `kernel-builder`.
e19323e verified
Raw
History Blame
9.66 kB
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import torch.distributed as dist
from ....ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd
from ....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 ....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 ....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
 )