Kernels:

kernels-community
/

fla

Trusted publisher

fla / build /torch-cuda /modules /token_shift_cp.py

Uploaded using `kernel-builder`.

e19323e verified 28 days ago

7.62 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

	"""
	Context Parallel support for Token Shift.

	Token shift has a 1-token dependency on previous tokens:
	y[t] = x[t-1] - x[t] (for t > 0)
	y[0] = cache - x[0] (cache is the last token from previous rank)

	In CP mode, non-first ranks need the last token from the previous rank as cache.
	Backward: non-last ranks need to send the last token's gradient to previous rank.
	"""

	import torch
	import torch.distributed as dist

	from ..modules.token_shift import token_shift_bwd, token_shift_fwd
	from ..ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd


	class TokenShiftCPFunction(torch.autograd.Function):
	"""
	Context Parallel version of TokenShift.

	Forward:
	1. Get last token from previous rank to construct cache
	2. Call token_shift_fwd with cache

	Backward:
	1. Call token_shift_bwd to get dx
	2. Sync communication: add next rank's first token gradient to current rank's last token
	"""

	@staticmethod
	def _prepare_cache_for_cp(
	x: torch.Tensor,
	cu_seqlens: torch.Tensor \| None,
	context: FLACPContext,
	group: dist.ProcessGroup \| None,
	) -> tuple[torch.Tensor \| None, int]:
	"""Prepare cache for CP forward pass by communicating with previous rank.

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

	Returns:
	cache: Cache tensor of shape [N, D] or None
	pre_num_tokens: Number of tokens from previous rank for the first sequence
	"""
	if group is None:
	return None, 0

	D = x.shape[-1]
	cache = None
	pre_num_tokens = 0

	if not context.is_first_rank:
	# Non-first rank: need cache from previous rank
	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]
	last_token = x_2d[-1:].contiguous() # [1, D]
	prev_last_token = conv_cp_send_recv_fwd(last_token, group) # [1, D]

	# For varlen: only the first sequence needs cache from prev rank
	N = len(cu_seqlens) - 1 if cu_seqlens is not None else 1
	cache = torch.zeros(N, D, device=x.device, dtype=x.dtype)

	# pre_num_conv_tokens tells us how many tokens from prev rank
	# belong to the first sequence on this rank
	pre_num_tokens = getattr(context, 'pre_num_conv_tokens', 0)
	if pre_num_tokens > 0:
	# The prev rank's last token is used as cache for first sequence
	cache[0] = prev_last_token[0]
	else:
	# First rank: participate in send but don't use received data
	x_2d = x.squeeze(0)
	last_token = x_2d[-1:].contiguous()
	_ = conv_cp_send_recv_fwd(last_token, group)

	return cache, pre_num_tokens

	@staticmethod
	def _correct_dx_for_cp(
	dx: torch.Tensor,
	grad_cache: torch.Tensor \| None,
	group: dist.ProcessGroup \| None,
	is_first_rank: bool,
	pre_num_tokens: int = 0,
	) -> None:
	"""Correct dx gradients for CP backward pass.

	Args:
	dx: Gradient tensor to be corrected, shape [1, T, D]
	grad_cache: Gradient w.r.t. cache, shape [N, D] or None
	group: Process group
	is_first_rank: Whether this is the first rank
	pre_num_tokens: Number of tokens from previous rank for first sequence
	"""
	if group is None:
	return

	D = dx.shape[-1]

	# Prepare gradient to send to previous rank
	if grad_cache is not None and pre_num_tokens > 0:
	# Only first sequence's cache gradient is relevant
	d_cache = grad_cache[0:1] # [1, D]
	else:
	d_cache = torch.zeros(1, D, device=dx.device, dtype=dx.dtype)

	# Send to previous rank, receive from next rank
	recv_grad = conv_cp_send_recv_bwd(d_cache, group) # [1, D]

	# Add received gradient to current rank's last token
	dx[0, -1, :].add_(recv_grad[0])

	@staticmethod
	def forward(
	ctx,
	x: torch.Tensor,
	cu_seqlens: torch.Tensor \| None,
	chunk_indices: torch.Tensor \| None,
	cp_context: FLACPContext \| None,
	):
	if cp_context is None:
	raise ValueError("cp_context must be provided for TokenShiftCPFunction")

	cu_seqlens = cp_context.cu_seqlens
	group = cp_context.group

	# Prepare cache for CP
	cache, pre_num_tokens = TokenShiftCPFunction._prepare_cache_for_cp(
	x=x,
	cu_seqlens=cu_seqlens,
	context=cp_context,
	group=group,
	)

	# Save for backward
	ctx.cu_seqlens = cu_seqlens
	ctx.chunk_indices = chunk_indices
	ctx.group = group
	ctx.has_cache = cache is not None
	ctx.is_first_rank = cp_context.is_first_rank
	ctx.pre_num_tokens = pre_num_tokens

	# Call original forward
	y, N, T, use_short_kernel, cache_out = token_shift_fwd(
	x=x,
	cu_seqlens=cu_seqlens,
	cache=cache,
	output_cache=True,
	chunk_indices=chunk_indices,
	)

	ctx.N = N
	ctx.T = T
	ctx.use_short_kernel = use_short_kernel

	return y

	@staticmethod
	def backward(ctx, dy: torch.Tensor):
	group = ctx.group

	# Prepare dcache for backward
	# For CP: non-last rank needs to receive gradient from next rank
	# This is handled in _correct_dx_for_cp after computing dx
	dcache = None # Will be computed by token_shift_bwd

	# Call original backward
	dx, grad_cache = token_shift_bwd(
	dy=dy,
	N=ctx.N,
	T=ctx.T,
	dcache=dcache,
	cu_seqlens=ctx.cu_seqlens,
	use_short_kernel=ctx.use_short_kernel,
	has_init_cache=ctx.has_cache,
	chunk_indices=ctx.chunk_indices,
	)

	# Correct dx gradients for CP
	TokenShiftCPFunction._correct_dx_for_cp(
	dx=dx,
	grad_cache=grad_cache,
	group=group,
	is_first_rank=ctx.is_first_rank,
	pre_num_tokens=ctx.pre_num_tokens,
	)

	return dx, None, None, None


	@torch.compiler.disable
	def token_shift_cp(
	x: torch.Tensor,
	cp_context: FLACPContext,
	cu_seqlens: torch.Tensor \| None = None,
	chunk_indices: torch.Tensor \| None = None,
	):
	"""
	Context Parallel version of token_shift.

	Args:
	x: Input tensor of shape [1, T, D]
	cp_context: CP context (required for CP mode)
	cu_seqlens: Cumulative sequence lengths
	chunk_indices: Chunk indices for variable-length sequences

	Returns:
	output: Tensor of shape [1, T, D] after applying token-shift
	"""
	if cp_context is None:
	raise ValueError("cp_context must be provided for token_shift_cp")

	assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for token_shift_cp"

	return TokenShiftCPFunction.apply(
	x, cu_seqlens, chunk_indices, cp_context
	)