diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index a62ca73187..1cd4ff1c25 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -56,6 +56,14 @@ # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" +_PACKED_CONTIGUOUS_ENV = "NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS" + + +def _use_packed_contiguous_thd(): + """Return whether experimental packed-contiguous THD all-gather is enabled.""" + # Read dynamically so input partitioning and attention agree when launchers or tests set + # the experimental flag after importing Transformer Engine. + return os.getenv(_PACKED_CONTIGUOUS_ENV, "0") == "1" def _reject_custom_recipe_under_cp(fp8, fp8_recipe): @@ -267,6 +275,129 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): return _seq_chunk_ids_cache_for_reordering_after_attn[(cp_size, device)] +def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None): + """Return physical boundaries used by the THD partition CUDA kernels.""" + target_device = torch.device(device if device is not None else cu_seqlens_padded.device) + target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype + return cu_seqlens_padded.to(device=target_device, dtype=target_dtype) + + +def _get_thd_partitioned_indices_reference(cu_seqlens_padded, total_tokens, cp_size, cp_rank): + """CPU fallback for dataloader-side THD partitioning.""" + total_chunks = 2 * cp_size + chunk_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) // total_chunks + indices = [] + for chunk_size, seq_start in zip(chunk_sizes, cu_seqlens_padded[:-1]): + indices.extend( + ( + torch.arange( + seq_start + cp_rank * chunk_size, + seq_start + (cp_rank + 1) * chunk_size, + ), + torch.arange( + seq_start + (total_chunks - cp_rank - 1) * chunk_size, + seq_start + (total_chunks - cp_rank) * chunk_size, + ), + ) + ) + return torch.cat(indices) + + +def get_thd_partitioned_indices( + cu_seqlens_padded, + total_tokens, + cp_size, + cp_rank, + device=None, +): + """Return THD token indices using the selected CP partition contract.""" + packed_contiguous = _use_packed_contiguous_thd() + if packed_contiguous: + validate_packed_contiguous_thd_metadata( + cu_seqlens_padded, + cu_seqlens_padded, + total_tokens, + cp_size, + ) + target_device = torch.device(device if device is not None else cu_seqlens_padded.device) + target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype + chunk_size = total_tokens // cp_size + return torch.arange( + cp_rank * chunk_size, + (cp_rank + 1) * chunk_size, + dtype=target_dtype, + device=target_device, + ) + cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, device) + if not cu_seqlens_padded.is_cuda: + return _get_thd_partitioned_indices_reference( + cu_seqlens_padded, total_tokens, cp_size, cp_rank + ) + if cu_seqlens_padded.dtype != torch.int32: + cu_seqlens_padded = cu_seqlens_padded.to(torch.int32) + return tex.thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_rank) + + +def validate_packed_contiguous_thd_metadata( + cu_seqlens, + cu_seqlens_padded, + total_tokens, + cp_size, +): + """Validate packed-contiguous THD metadata while producing rank-local inputs.""" + assert cu_seqlens.shape == cu_seqlens_padded.shape + assert total_tokens % cp_size == 0 + assert cu_seqlens[0] == 0 and cu_seqlens_padded[0] == 0 + assert cu_seqlens_padded[-1] == total_tokens + assert torch.all(cu_seqlens[1:] >= cu_seqlens[:-1]) + assert torch.all(cu_seqlens_padded[1:] >= cu_seqlens_padded[:-1]) + actual_seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + padded_seqlens = cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] + assert torch.all(actual_seqlens <= padded_seqlens) + + +def get_packed_contiguous_thd_causal_metadata( + cu_seqlens, + cu_seqlens_padded, + total_tokens, + cp_size, + cp_rank, +): + """Build one-step THD metadata for packed-contiguous CP partitioning. + + The complete physical token buffer is the sharding unit. ``cu_seqlens`` + remains the logical document boundary, so each global chunk is represented + as its intersection with every document. + """ + assert cp_size > 0 and 0 <= cp_rank < cp_size + assert cu_seqlens.shape == cu_seqlens_padded.shape + assert total_tokens % cp_size == 0 + + chunk_size = total_tokens // cp_size + chunk_start = cp_rank * chunk_size + chunk_end = chunk_start + chunk_size + actual_seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + doc_starts = cu_seqlens_padded[:-1] + valid_doc_ends = doc_starts + actual_seqlens + + fragment_starts = torch.clamp(doc_starts, min=chunk_start, max=chunk_end) + fragment_ends = torch.clamp(valid_doc_ends, min=chunk_start, max=chunk_end) + fragment_seqlens = torch.clamp(fragment_ends - fragment_starts, min=0) + + q_cu_seqlens = torch.zeros_like(cu_seqlens) + q_cu_seqlens[1:] = fragment_seqlens.cumsum(0) + q_cu_seqlens_padded = ( + torch.clamp(cu_seqlens_padded, min=chunk_start, max=chunk_end) - chunk_start + ) + + visible_kv_seqlens = torch.clamp(chunk_end - doc_starts, min=0) + visible_kv_seqlens = torch.minimum(visible_kv_seqlens, actual_seqlens) + kv_cu_seqlens = torch.zeros_like(cu_seqlens) + kv_cu_seqlens[1:] = visible_kv_seqlens.cumsum(0) + + return [q_cu_seqlens], [q_cu_seqlens_padded], [kv_cu_seqlens] + + @jit_fuser def reorder_seq_chunks_for_a2a_before_attn(x, chunk_ids_for_a2a, seq_dim, cp_size): """Reorder sequence chunk for A2A communication before attention compute.""" @@ -384,6 +515,24 @@ def thd_cp_rank_order_to_sequence_order(x, cu_seqlens, cp_size, seq_dim=0): return tex.thd_cp_rank_order_to_sequence_order(x, cu_seqlens, cp_size, x.shape[seq_dim]) +def restore_thd_gathered_kv(x, cu_seqlens_padded, cp_size): + """Restore gathered THD tokens to physical sequence order.""" + if _use_packed_contiguous_thd(): + # Rank r owns physical chunk r, so rank-major all-gather is already in sequence order. + return x + cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, x.device) + return thd_cp_rank_order_to_sequence_order(x, cu_seqlens_padded, cp_size) + + +def unrestore_thd_gathered_kv(x, cu_seqlens_padded, cp_size): + """Arrange physical THD tokens for rank-ordered reduce-scatter.""" + if _use_packed_contiguous_thd(): + # Physical sequence order is also the rank-major reduce-scatter order for this policy. + return x + cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, x.device) + return thd_sequence_order_to_cp_rank_order(x, cu_seqlens_padded, cp_size) + + def flash_attn_a2a_communicate( a2a_inputs: Union[torch.Tensor, List[torch.Tensor]], chunk_ids_for_a2a: torch.Tensor, @@ -3073,6 +3222,29 @@ def forward( if qkv_format == "thd": # THD always uses padding mask types; per-step masks set internally assert padding, f"THD format requires padding mask type, got {attn_mask_type}!" + packed_contiguous = qkv_format == "thd" and _use_packed_contiguous_thd() + if packed_contiguous: + assert qkv_format == "thd" + assert ( + use_fused_attention or use_flash_attn_3 + ), "Packed THD partitioning requires FusedAttention or FlashAttention 3." + assert not ( + use_flash_attn_3 and pad_between_seqs + ), "Packed THD partitioning with FlashAttention 3 does not support padding yet." + assert causal and window_size == ( + -1, + 0, + ), "Packed THD partitioning currently supports full causal attention only." + assert not fp8, "Packed THD partitioning does not support FP8 yet." + assert ( + not is_graph_capturing() + ), "Packed THD partitioning does not support CUDA graph capture yet." + assert ( + q.shape[0] == k.shape[0] == v.shape[0] + ), "Packed THD partitioning requires equal local Q/K/V physical lengths." + assert cu_seqlens_q is cu_seqlens_kv and ( + cu_seqlens_q_padded is cu_seqlens_kv_padded + ), "Packed THD self-attention requires shared Q/KV sequence metadata tensors." # AG CP uses shorter per-step Q against longer KV, so causal masks need # bottom-right alignment for both sliced and THD paths. if use_fused_attention and causal and "bottom_right" not in attn_mask_type: @@ -3092,10 +3264,11 @@ def forward( "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." ) - assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( - "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" - f" {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}." - ) + if not packed_contiguous: + assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( + "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found " + f"seq_len_q = {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}." + ) flash_attn_fwd = None if not use_fused_attention: @@ -3140,14 +3313,18 @@ def forward( q.shape[seq_dim] % 2 == 0 and k.shape[seq_dim] % 2 == 0 ), "Sequence length per GPU needs to be divisible by 2!" - # Divide by 2*cp_size to get per-chunk values - max_seqlen_q = max_seqlen_q // (2 * cp_size) - max_seqlen_kv = max_seqlen_kv // (2 * cp_size) + # Per-document DCS divides every sequence into 2*CP chunks. Packed-contiguous + # instead bounds Q by one global chunk and keeps full-document KV bounds. + if packed_contiguous: + max_seqlen_q = min(max_seqlen_q, q.shape[0]) + else: + max_seqlen_q = max_seqlen_q // (2 * cp_size) + max_seqlen_kv = max_seqlen_kv // (2 * cp_size) if use_fused_attention and qkv_format != "thd": cu_seqlens_q = cu_seqlens_q // (2 * cp_size) - if qkv_format == "thd": + if qkv_format == "thd" and not packed_contiguous: cu_seqlens_q_padded = cu_seqlens_q_padded // (2 * cp_size) - else: + elif qkv_format != "thd": cu_seqlens_q_padded = None if use_fused_attention and attn_mask_type == "causal": attn_mask_type = attn_mask_type + "_bottom_right" @@ -3215,10 +3392,8 @@ def forward( if qkv_format == "thd": # [cp*t, h, d] -> reorder to sequence order -> [t_full, h, d] - # The padded cu_seqlens are global sequence offsets. Reorder uses them to - # derive per-sequence chunk boundaries. - k_ag = thd_cp_rank_order_to_sequence_order(k_ag, cu_seqlens_kv_padded, cp_size) - v_ag = thd_cp_rank_order_to_sequence_order(v_ag, cu_seqlens_kv_padded, cp_size) + k_ag = restore_thd_gathered_kv(k_ag, cu_seqlens_kv_padded, cp_size) + v_ag = restore_thd_gathered_kv(v_ag, cu_seqlens_kv_padded, cp_size) else: # [cp, s, b, h, d] -> [cp*2, s//2, b, h, d] k_ag = k_ag.view(2 * cp_size, k.shape[0] // 2, *k.shape[1:]) @@ -3249,7 +3424,7 @@ def forward( # create two streams to resolve wave quantization issue of Flash Attn in each step flash_attn_streams = [torch.cuda.current_stream(), cp_stream] # prepare per-step tensors - local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] + local_seq_chunk_ids = [rank] if packed_contiguous else [rank, 2 * cp_size - rank - 1] kv_seq_range_per_step = [None, None] window_size_per_step = [None, None] cu_seqlens_kv_per_step = [None, None] @@ -3263,7 +3438,20 @@ def forward( max_logit = None # Pre-compute THD-specific per-step cu_seqlens - if qkv_format == "thd": + if qkv_format == "thd" and packed_contiguous: + total_tokens_q = q.shape[0] * cp_size + ( + thd_cu_seqlens_q_per_step, + thd_cu_seqlens_q_padded_per_step, + thd_cu_seqlens_kv_per_step, + ) = get_packed_contiguous_thd_causal_metadata( + cu_seqlens_q_original, + cu_seqlens_q_padded, + total_tokens_q, + cp_size, + rank, + ) + elif qkv_format == "thd": # Rank-level padded offsets (2 chunks per sequence on this rank) cu_seqlens_q_padded_rank = cu_seqlens_q_padded * 2 @@ -3334,6 +3522,11 @@ def forward( thd_cu_seqlens_kv_per_step[0][1:] = visible_actual[0].cumsum(0) thd_cu_seqlens_kv_per_step[1][1:] = visible_actual[1].cumsum(0) + # Step 1 runs on cp_stream and consumes THD metadata produced above on + # the current stream. The earlier wait only covered K/V AG and reorder. + if qkv_format == "thd": + cp_stream.wait_stream(torch.cuda.current_stream()) + for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): # FA3 uses internal per-call workspace. Consecutive AG per-step @@ -3393,15 +3586,19 @@ def forward( q_part = q k_part = k_ag v_part = v_ag - kv_range, window_size_per_step[i] = get_kv_seq_info_after_all_gather( - local_seq_chunk_ids[i], - cp_size, - max_seqlen_q, - max_seqlen_kv, - window_size, - causal, - ) - max_seqlen_kv_ = kv_range[1] + if packed_contiguous: + window_size_per_step[i] = (-1, 0) + max_seqlen_kv_ = max_seqlen_kv + else: + kv_range, window_size_per_step[i] = get_kv_seq_info_after_all_gather( + local_seq_chunk_ids[i], + cp_size, + max_seqlen_q, + max_seqlen_kv, + window_size, + causal, + ) + max_seqlen_kv_ = kv_range[1] cu_seqlens_kv_per_step[i] = thd_cu_seqlens_kv_per_step[i] if use_fused_attention: # Set per-step parameters for THD vs bshd/sbhd @@ -3646,6 +3843,7 @@ def forward( ctx.use_flash_attn_3 = use_flash_attn_3 ctx.pad_between_seqs = pad_between_seqs ctx.window_size = window_size + ctx.packed_contiguous = packed_contiguous if qkv_format == "thd": ctx.max_seqlen_kv = max_seqlen_kv ctx.cu_seqlens_kv_padded = cu_seqlens_kv_padded @@ -3790,9 +3988,8 @@ def backward(ctx, dout, *_args): cu_seqlens_kv_padded = ctx.cu_seqlens_kv_padded thd_cu_seqlens_q_per_step = ctx.thd_cu_seqlens_q_per_step # [cp*t, h, d] -> reorder to sequence order - # Use padded cu_seqlens (divisible by 2*cp_size) for correct reorder - k_ag = thd_cp_rank_order_to_sequence_order(k_ag, cu_seqlens_kv_padded, cp_size) - v_ag = thd_cp_rank_order_to_sequence_order(v_ag, cu_seqlens_kv_padded, cp_size) + k_ag = restore_thd_gathered_kv(k_ag, cu_seqlens_kv_padded, cp_size) + v_ag = restore_thd_gathered_kv(v_ag, cu_seqlens_kv_padded, cp_size) thd_cu_seqlens_q_padded_per_step = ctx.thd_cu_seqlens_q_padded_per_step else: @@ -3840,7 +4037,7 @@ def backward(ctx, dout, *_args): if fa_utils.v2_6_0_plus: fa_backward_kwargs["softcap"] = 0.0 - local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] + local_seq_chunk_ids = [rank] if ctx.packed_contiguous else [rank, 2 * cp_size - rank - 1] for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): # FA3 uses internal per-call workspace. Consecutive AG per-step @@ -3855,15 +4052,18 @@ def backward(ctx, dout, *_args): q_part = q k_part = k_ag v_part = v_ag - kv_range, _ = get_kv_seq_info_after_all_gather( - local_seq_chunk_ids[i], - cp_size, - ctx.max_seqlen_q, - ctx.max_seqlen_kv, - ctx.window_size, - "causal" in ctx.attn_mask_type, - ) - max_seqlen_kv = kv_range[1] + if ctx.packed_contiguous: + max_seqlen_kv = ctx.max_seqlen_kv + else: + kv_range, _ = get_kv_seq_info_after_all_gather( + local_seq_chunk_ids[i], + cp_size, + ctx.max_seqlen_q, + ctx.max_seqlen_kv, + ctx.window_size, + "causal" in ctx.attn_mask_type, + ) + max_seqlen_kv = kv_range[1] out_part = out dout_part = dout else: @@ -4099,9 +4299,8 @@ def backward(ctx, dout, *_args): if ctx.qkv_format == "thd": # Reorder dK/dV from sequence order back to dual-chunk CP rank order, # then reduce-scatter across CP ranks. - # Use padded cu_seqlens for correct slice boundaries. - dk = thd_sequence_order_to_cp_rank_order(dk, cu_seqlens_kv_padded, cp_size) - dv = thd_sequence_order_to_cp_rank_order(dv, cu_seqlens_kv_padded, cp_size) + dk = unrestore_thd_gathered_kv(dk, cu_seqlens_kv_padded, cp_size) + dv = unrestore_thd_gathered_kv(dv, cu_seqlens_kv_padded, cp_size) dk, _ = reduce_scatter_along_first_dim(dk, ctx.cp_group) dv, _ = reduce_scatter_along_first_dim(dv, ctx.cp_group) # dQ is already [t_rank, h, d], no reshape needed @@ -4956,10 +5155,18 @@ def attn_forward_func_with_cp( """ Attention implementation with context parallelism (CP). CP partitions tensors along the sequence dimension, and by reducing the memory and computational pressure on each GPU, it enables long-context - LLMs in a distributed fashion. Transformer Engine's PyTorch CP implementation currently utilizes - the DualChunkSwap strategy to ensure load balancing across CP ranks. It is applied to all `attn_mask_type`s - and all `qkv_format`s, and it requires sequence lengths to be, or are padded to be, divisible by - (cp_size * 2). It also requires tokens to be re-ordered before entering this function. + LLMs in a distributed fashion. By default, Transformer Engine's PyTorch CP + implementation applies DualChunkSwap to each sequence independently. It requires + every sequence length to be, or be padded to be, divisible by (cp_size * 2), and + tokens must be re-ordered before entering this function. + + Experimental environment flag + ``NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS=1`` instead assigns one + contiguous physical-buffer chunk to each rank and uses one attention step per rank. + Logical sequences remain isolated by ``cu_seqlens``. This mode requires THD, + all-gather, full causal self-attention, and FusedAttention, or FlashAttention 3 + without padding. Input producers must use the same flag when partitioning inputs + with :func:`get_batch_on_this_cp_rank` or :func:`get_thd_partitioned_indices`. For qkv_format = {'bshd', 'sbhd'}, the token re-ordering is illustrated as below, for an example use case of s = 12, attn_mask_type = 'causal', and cp_size = 2. seq_pos indicates each token's position @@ -5040,6 +5247,11 @@ def attn_forward_func_with_cp( "sbhd", "thd", ], f"Context parallelism does not support {qkv_format=}!" + packed_contiguous = qkv_format == "thd" and _use_packed_contiguous_thd() + if packed_contiguous: + assert ( + qkv_format == "thd" and cp_comm_type == "all_gather" + ), "Packed THD partitioning requires qkv_format='thd' and cp_comm_type='all_gather'." assert ( qkv_format != "sbhd" or use_fused_attention ), "Context parallelism does not support FlashAttention backend with qkv_format = 'sbhd'!" @@ -5269,6 +5481,8 @@ def get_batch_on_this_cp_rank( Which are parallelized across GPUs in a context parallel group. This version works with variable-length sequences using cumulative sequence lengths. + ``NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS=1`` assigns one contiguous + physical-buffer chunk per rank. By default, each padded sequence is chunked independently. """ if qvk_format not in ["thd", "bshd", "sbhd"]: raise ValueError(f"Unsupported qvk_format: {qvk_format}!") @@ -5278,12 +5492,6 @@ def get_batch_on_this_cp_rank( if cp_size > 1: cp_rank = torch.distributed.get_rank(group=cp_group) - # Calculate the chunk sizes for each sequence - total_slices_of_any_sequence = 2 * cp_size - slice_sizes = ( - cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] - ) // total_slices_of_any_sequence - # Process each tensor directly instead of using keys_to_change loop def process_tensor(val): if val is None: @@ -5316,29 +5524,14 @@ def process_tensor(val): else: raise ValueError("Tensor must be at least 1D") - # On this particular rank, for each sequence, get two slices, one from the beginning - # and one from the end. - cp_rank_slices = [] - for slice_size, seq_start in zip(slice_sizes, cu_seqlens_padded[:-1]): - # 1st segment - cp_rank_slices.append( - torch.arange( - seq_start + (cp_rank * slice_size), - seq_start + ((cp_rank + 1) * slice_size), - device=val.device, - ) - ) - - # 2nd segment - cp_rank_slices.append( - torch.arange( - seq_start + ((total_slices_of_any_sequence - cp_rank - 1) * slice_size), - seq_start + ((total_slices_of_any_sequence - cp_rank) * slice_size), - device=val.device, - ) - ) - - return val.index_select(current_seq_dim, torch.cat(cp_rank_slices)) + cp_rank_indices = get_thd_partitioned_indices( + cu_seqlens_padded, + seq_len_val, + cp_size, + cp_rank, + val.device, + ) + return val.index_select(current_seq_dim, cp_rank_indices) # Process each tensor directly input_ids_padded = process_tensor(input_ids_padded)