`core.context_parallel.layout`#

Context-parallel sequence-layout conversion.

Module Contents#

Classes#

`_LayoutRedistributionPlan`	Rank-local all-to-all plan expressed in sequence-segment counts.
`THDCPLayoutPlan`	Rank-local all-to-all-v plan for packed sequence layout conversion.
`_LayoutParallelContext`	Process-group coordinates used by a layout redistribution.
`ContextParallelLayoutManager`	Manage CP layout transitions across a sequence of layers.
`ContextParallelLayoutState`	Per-forward state for CP layout conversion.

Functions#

`_segments_per_rank`	Return two segments for CP-only conversion and one for even-TP SP conversion.
`_local_segment_ids`	Return the atomic sequence segments owned by one TP x CP rank.
`_segment_owner`
`_build_group_rank_by_logical_rank`	Map logical `cp_rank * tp_size + tp_rank` coordinates to group ranks.
`_get_group_rank_by_logical_rank`
`_get_layout_parallel_context`	Resolve the rank coordinates and process group for a layout redistribution.
`_build_layout_redistribution_plan`	Build the all-to-all-v plan for one rank of a CP layout conversion.
`_build_thd_cp_layout_plan_from_rank_order_indices`	Build a packed layout plan from source indices grouped by target logical rank.
`_build_thd_rank_order_indices`	Build padded dual-chunk attention order without Transformer Engine helpers.
`build_thd_cp_layout_plan`	Build a reusable packed sequence layout plan for attention.
`_redistribute_layout`	Redistribute local sequence segments with a differentiable all-to-all-v.
`_redistribute_thd_layout`	Apply a packed sequence layout plan with a differentiable all-to-all-v.
`contiguous_to_zigzag`	Convert contiguous CP sequence shards to Megatron’s zigzag attention layout.
`zigzag_to_contiguous`	Convert Megatron’s zigzag attention shards back to contiguous CP shards.

Data#

CPLayout

API#

core.context_parallel.layout.CPLayout#: None

class core.context_parallel.layout._LayoutRedistributionPlan#

Rank-local all-to-all plan expressed in sequence-segment counts.

send_slots: tuple[int, ...]#: None

input_segment_counts: tuple[int, ...]#: None

output_segment_counts: tuple[int, ...]#: None

receive_permutation: tuple[int, ...]#: None

class core.context_parallel.layout.THDCPLayoutPlan#

Rank-local all-to-all-v plan for packed sequence layout conversion.

Forward fields describe contiguous-to-zigzag conversion. Reverse fields describe the inverse conversion. The plan is built once per packed batch and reused around every attention layer.

contiguous_local_token_count: int#: None

zigzag_local_token_count: int#: None

cu_seqlens_padded: torch.Tensor#: None

max_seqlen_padded: int#: None

pad_between_seqs: bool#: None

forward_send_indices: torch.Tensor#: None

forward_receive_positions: torch.Tensor#: None

forward_input_split_sizes: tuple[int, ...]#: None

forward_output_split_sizes: tuple[int, ...]#: None

reverse_send_indices: torch.Tensor#: None

reverse_receive_indices: torch.Tensor#: None

class core.context_parallel.layout._LayoutParallelContext#

Process-group coordinates used by a layout redistribution.

cp_size: int#: None

cp_rank: int#: None

tp_size: int#: None

tp_rank: int#: None

communication_group: torch.distributed.ProcessGroup#: None

group_rank_by_logical_rank: tuple[int, ...]#: None

property group_size: int#: Return the number of ranks participating in the redistribution.

core.context_parallel.layout._segments_per_rank(tp_size: int) → int#: Return two segments for CP-only conversion and one for even-TP SP conversion.

core.context_parallel.layout._local_segment_ids( layout: core.context_parallel.layout.CPLayout, cp_size: int, cp_rank: int, tp_size: int = 1, tp_rank: int = 0, ) → tuple[int, ...]#: Return the atomic sequence segments owned by one TP x CP rank.

core.context_parallel.layout._segment_owner( segment_id: int, layout: core.context_parallel.layout.CPLayout, cp_size: int, tp_size: int, ) → tuple[int, int]#

core.context_parallel.layout._build_group_rank_by_logical_rank( cp_global_ranks: tuple[int, ...], tp_global_ranks: tuple[int, ...], tp_cp_global_ranks: tuple[int, ...], current_global_rank: int, ) → tuple[int, ...]#: Map logical cp_rank * tp_size + tp_rank coordinates to group ranks.

core.context_parallel.layout._get_group_rank_by_logical_rank( cp_group: torch.distributed.ProcessGroup, tp_group: torch.distributed.ProcessGroup, tp_cp_group: torch.distributed.ProcessGroup, ) → tuple[int, ...]#

core.context_parallel.layout._get_layout_parallel_context( cp_group: torch.distributed.ProcessGroup, sequence_parallel: bool, tp_group: torch.distributed.ProcessGroup | None, tp_cp_group: torch.distributed.ProcessGroup | None, ) → core.context_parallel.layout._LayoutParallelContext#: Resolve the rank coordinates and process group for a layout redistribution.

core.context_parallel.layout._build_layout_redistribution_plan( source_layout: core.context_parallel.layout.CPLayout, target_layout: core.context_parallel.layout.CPLayout, cp_size: int, cp_rank: int, tp_size: int = 1, tp_rank: int = 0, group_rank_by_logical_rank: tuple[int, ...] | None = None, ) → core.context_parallel.layout._LayoutRedistributionPlan#: Build the all-to-all-v plan for one rank of a CP layout conversion.

core.context_parallel.layout._build_thd_cp_layout_plan_from_rank_order_indices( rank_order_indices: torch.Tensor, source_token_count: int, cu_seqlens_padded: torch.Tensor, cp_size: int, cp_rank: int, tp_size: int = 1, tp_rank: int = 0, group_rank_by_logical_rank: tuple[int, ...] | None = None, pad_between_seqs: bool | torch.Tensor = False, ) → core.context_parallel.layout.THDCPLayoutPlan#

Build a packed layout plan from source indices grouped by target logical rank.

rank_order_indices contains one entry for every padded attention position. Non-negative entries identify positions in the contiguous input; -1 entries are temporary padding.

core.context_parallel.layout._build_thd_rank_order_indices( cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor | None, cp_size: int, tp_size: int, expected_source_token_count: int | None = None, ) → tuple[torch.Tensor, torch.Tensor]#: Build padded dual-chunk attention order without Transformer Engine helpers.

core.context_parallel.layout.build_thd_cp_layout_plan( cu_seqlens: torch.Tensor, total_tokens: int, cp_group: torch.distributed.ProcessGroup, cu_seqlens_padded: torch.Tensor | None = None, sequence_parallel: bool = False, tp_group: torch.distributed.ProcessGroup | None = None, tp_cp_group: torch.distributed.ProcessGroup | None = None, ) → core.context_parallel.layout.THDCPLayoutPlan#

Build a reusable packed sequence layout plan for attention.

Parameters:

cu_seqlens – Global cumulative actual sequence lengths in torch.int32.
total_tokens – Global token count in the contiguous residual stream.
cp_group – Context-parallel process group.
cu_seqlens_padded – Optional physical offsets already present in the input.
sequence_parallel – Whether the residual stream is also sharded over TP ranks.
tp_group – Tensor-parallel process group, required with sequence parallelism.
tp_cp_group – Combined TP x CP process group, required when TP size is greater than one.

Returns:

A rank-local plan reusable for both directions of the layout conversion.

core.context_parallel.layout._redistribute_layout( input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup, source_layout: core.context_parallel.layout.CPLayout, target_layout: core.context_parallel.layout.CPLayout, sequence_parallel: bool, tp_group: torch.distributed.ProcessGroup | None, tp_cp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#: Redistribute local sequence segments with a differentiable all-to-all-v.

core.context_parallel.layout._redistribute_thd_layout( input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup, plan: core.context_parallel.layout.THDCPLayoutPlan, source_layout: core.context_parallel.layout.CPLayout, target_layout: core.context_parallel.layout.CPLayout, sequence_parallel: bool, tp_group: torch.distributed.ProcessGroup | None, tp_cp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#: Apply a packed sequence layout plan with a differentiable all-to-all-v.

core.context_parallel.layout.contiguous_to_zigzag( input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup, sequence_parallel: bool = False, tp_group: torch.distributed.ProcessGroup | None = None, tp_cp_group: torch.distributed.ProcessGroup | None = None, thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None = None, ) → torch.Tensor#: Convert contiguous CP sequence shards to Megatron’s zigzag attention layout.

core.context_parallel.layout.zigzag_to_contiguous( input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup, sequence_parallel: bool = False, tp_group: torch.distributed.ProcessGroup | None = None, tp_cp_group: torch.distributed.ProcessGroup | None = None, thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None = None, ) → torch.Tensor#: Convert Megatron’s zigzag attention shards back to contiguous CP shards.

class core.context_parallel.layout.ContextParallelLayoutManager#

Manage CP layout transitions across a sequence of layers.

layer_layouts: tuple[core.context_parallel.layout.CPLayout, ...]#: None

boundary_layout: core.context_parallel.layout.CPLayout#: None

sequence_parallel: bool#: None

cp_group: torch.distributed.ProcessGroup#: None

tp_group: torch.distributed.ProcessGroup | None#: None

tp_cp_group: torch.distributed.ProcessGroup | None#: None

requires_conversion: bool#: ‘field(…)’

__post_init__() → None#: Determine whether the layer sequence needs layout conversion.

_convert_cp_layout( hidden_states: torch.Tensor, source_layout: core.context_parallel.layout.CPLayout, target_layout: core.context_parallel.layout.CPLayout, thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None = None, ) → torch.Tensor#

prepare_layer_input( layer_index: int, hidden_states: torch.Tensor, thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None = None, ) → torch.Tensor#: Convert when a layer requires a different layout than its predecessor.

finalize_layer_output( layer_index: int, hidden_states: torch.Tensor, thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None = None, ) → torch.Tensor#: Restore the boundary layout after the final layer.

build_packed_zigzag_layout( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams, ) → tuple[core.context_parallel.layout.THDCPLayoutPlan, megatron.core.packed_seq_params.PackedSeqParams]#: Build one THD conversion plan and its zigzag-layout metadata.

build_forward_state( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, ) → ContextParallelLayoutState | None#: Build the layout state for one forward pass.

class core.context_parallel.layout.ContextParallelLayoutState#

Per-forward state for CP layout conversion.

manager: core.context_parallel.layout.ContextParallelLayoutManager#: None

thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None#: None

contiguous_packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None#: None

zigzag_packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None#: None

prepare_layer( layer_index: int, hidden_states: torch.Tensor, ) → tuple[torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams | None]#: Prepare a layer’s input and matching packed metadata.

finalize_layer( layer_index: int, hidden_states: torch.Tensor, ) → torch.Tensor#: Finalize a layer’s output layout.

core.context_parallel.layout#