Kernels:
Trusted publisher
File size: 9,658 Bytes
e19323e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 | # 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
)
|