Kernels:

kernels-community
/

fla

Trusted publisher

fla / build /torch-cuda /ops /cp /context.py

Uploaded using `kernel-builder`.

e19323e verified 28 days ago

7.06 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

	from __future__ import annotations

	from dataclasses import dataclass
	from typing import TYPE_CHECKING

	import torch
	import torch.distributed as dist

	from ...utils import tensor_cache

	if TYPE_CHECKING:
	from torch.distributed import ProcessGroup


	@dataclass
	class FLACPContext:
	"""FLA Context Parallel Context - Operator-level context management."""
	group: ProcessGroup \| None = None
	cu_seqlens: torch.Tensor \| None = None
	cu_seqlens_cpu: torch.Tensor \| None = None
	is_last_rank: bool \| None = None
	pre_num_ranks: int \| None = None
	is_first_rank: bool \| None = None
	post_num_ranks: int \| None = None
	conv1d_kernel_size: int \| None = None
	pre_num_conv_tokens: int \| None = None

	def copy_for_backward(self) -> FLACPContext:
	"""Create a copy for backward pass (useful when PP_SIZE > 1)."""
	return FLACPContext(
	group=self.group,
	cu_seqlens=self.cu_seqlens.clone() if self.cu_seqlens is not None else None,
	cu_seqlens_cpu=self.cu_seqlens_cpu.clone() if self.cu_seqlens_cpu is not None else None,
	is_last_rank=self.is_last_rank,
	pre_num_ranks=self.pre_num_ranks,
	is_first_rank=self.is_first_rank,
	post_num_ranks=self.post_num_ranks,
	conv1d_kernel_size=self.conv1d_kernel_size,
	pre_num_conv_tokens=self.pre_num_conv_tokens,
	)

	@property
	def num_seqs(self) -> int:
	"""Number of sequences in this rank."""
	return 0 if self.cu_seqlens is None else len(self.cu_seqlens) - 1

	@property
	def is_cp_enabled(self) -> bool:
	"""Whether context parallel is enabled."""
	return self.group is not None


	@tensor_cache
	def get_cp_cu_seqlens(
	cu_seqlens: torch.LongTensor,
	cu_seqlens_cpu: torch.LongTensor \| None = None,
	world_size: int \| None = None,
	rank: int \| None = None,
	group: dist.ProcessGroup \| None = None,
	conv1d_kernel_size: int \| None = None
	) -> FLACPContext:
	# 1. Initialize environment info
	if world_size is None:
	assert group is not None
	world_size = dist.get_world_size(group=group)
	rank = dist.get_rank(group=group)

	# 2. Operate on CPU to avoid D2H sync and leverage vectorization (int64/long)
	if cu_seqlens_cpu is None:
	cu_seqlens_cpu = cu_seqlens.cpu()
	cu_seqlens_cpu = cu_seqlens_cpu.to(dtype=torch.long)

	# Get total tokens and current rank's responsible range
	# Assume cu_seqlens is [0, s1, s1+s2, ..., total]
	total_tokens = cu_seqlens_cpu[-1].item()
	part_len = total_tokens // world_size
	rank_start = part_len * rank
	rank_end = rank_start + part_len

	# 3. Vectorized search: find sequences overlapping with current rank's interval [rank_start, rank_end)
	# We need to find idx such that: global_ends[idx] > rank_start AND global_starts[idx] < rank_end

	# Optimization: cu_seqlens is sorted, use searchsorted to quickly locate boundaries
	# Find first sequence whose end > rank_start
	# cu_seqlens_cpu[1:] contains all sequence end points
	start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], rank_start, side='right')

	# Find first sequence whose start >= rank_end, sequences before this may overlap
	# cu_seqlens_cpu[:-1] contains all sequence start points
	end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], rank_end, side='left')

	# Slice cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] to get relevant global cu_seqlens nodes
	# +1 because end_seq_idx is an open boundary, and cu_seqlens length is num_seqs + 1
	subset_cu_seqlens = cu_seqlens_cpu[start_seq_idx: end_seq_idx + 1]

	# 4. Compute local cu_seqlens on CPU (int32)
	# Clamp global coordinates to [rank_start, rank_end], subtract rank_start to get local coordinates
	# unique_consecutive removes duplicates from clamping (e.g., sequences entirely outside this rank)
	local_cu_seqlens_cpu = (
	subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start
	).unique_consecutive().to(torch.int32)

	# Transfer to GPU (int32, small tensor, fast transfer)
	# non_blocking=True can further hide latency in CUDA streams
	local_cu_seqlens_gpu = local_cu_seqlens_cpu.to(
	device=cu_seqlens.device, non_blocking=True
	)

	# 5. Compute Context Parallel metadata (first/last rank info)
	# Use slice endpoints directly, avoiding loops

	# Get global info for the first sequence that has data on current rank
	first_seq_global_start = cu_seqlens_cpu[start_seq_idx].item()
	# Get global info for the last sequence that has data on current rank
	last_seq_global_end = cu_seqlens_cpu[end_seq_idx].item()

	# Number of tokens current rank needs from previous ranks for conv
	pre_num_conv_tokens = max(0, rank_start - first_seq_global_start)

	# Compute first sequence's starting rank
	first_rank_of_first_seq = first_seq_global_start // part_len
	# Number of previous ranks current rank needs to receive state from
	pre_num_ranks = rank - first_rank_of_first_seq
	# Whether current rank is the first in the sequence's processing chain
	is_first_rank = (rank == first_rank_of_first_seq)

	# Compute last sequence's ending rank
	# (last_seq_global_end - 1) is the index of the last token
	last_rank_of_last_seq = (last_seq_global_end - 1) // part_len
	# Number of subsequent ranks current rank needs to send state to
	post_num_ranks = last_rank_of_last_seq - rank
	# Whether current rank is the last in the sequence's processing chain
	is_last_rank = (rank == last_rank_of_last_seq)

	return FLACPContext(
	group=group,
	cu_seqlens=local_cu_seqlens_gpu,
	cu_seqlens_cpu=local_cu_seqlens_cpu,
	is_last_rank=is_last_rank,
	pre_num_ranks=pre_num_ranks,
	is_first_rank=is_first_rank,
	post_num_ranks=post_num_ranks,
	conv1d_kernel_size=conv1d_kernel_size,
	pre_num_conv_tokens=pre_num_conv_tokens
	)


	def build_cp_context(
	cu_seqlens: torch.Tensor,
	group: ProcessGroup,
	conv1d_kernel_size: int \| None = None,
	cu_seqlens_cpu: torch.Tensor \| None = None,
	) -> FLACPContext:
	"""Build a CP context for the given cu_seqlens and process group.

	Args:
	cu_seqlens: Cumulative sequence lengths tensor (before partition).
	group: Process group for CP communication.
	conv1d_kernel_size: Kernel size for convolution (optional).
	cu_seqlens_cpu: CPU version of cu_seqlens to avoid d2h transfer (optional).

	Returns:
	FLACPContext with computed cu_seqlens and rank information.
	"""
	return get_cp_cu_seqlens(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, group=group, conv1d_kernel_size=conv1d_kernel_size)