core.context_parallel.utils

Context-parallel batch partitioning helpers.

Module Contents

Classes

ContextParallelBatch

Batch views and reusable metadata prepared for the requested CP layouts.

Functions

_get_batch_on_this_cp_rank_contiguous	Use contiguous CP shards while keeping dense attention-mask queries zigzag.
_get_batch_on_this_cp_rank_padded_zigzag	Build a packed zigzag shard padded for TE context-parallel attention.
_build_packed_seq_params	Build packed sequence metadata for one physical CP layout.
get_batches_on_this_cp_rank	Partition a batch and prepare metadata for the requested CP layouts.

API

class core.context_parallel.utils.ContextParallelBatch

Batch views and reusable metadata prepared for the requested CP layouts.

Entries with the same layout key in batches_by_layout and packed_seq_params_by_layout describe the same token ordering. Dense attention masks are the exception: their query rows remain zigzag because that is the layout consumed by softmax attention. Accessors default to boundary_layout, and thd_plan connects the two packed token orderings.

boundary_layout: core.context_parallel.layout.CPLayout

None

batches_by_layout: dict[core.context_parallel.layout.CPLayout, Dict[str, Any]]

None

packed_seq_params_by_layout: dict[core.context_parallel.layout.CPLayout, megatron.core.packed_seq_params.PackedSeqParams | None]

None

thd_plan: core.context_parallel.layout.THDCPLayoutPlan | None

None

classmethod from_single_layout(

layout: core.context_parallel.layout.CPLayout,

batch: Dict[str, Any],

packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None,

) → core.context_parallel.utils.ContextParallelBatch

Wrap an already-partitioned batch that has one physical CP layout.

get_batch(

layout: core.context_parallel.layout.CPLayout | None = None,

) → Dict[str, Any]

Return the batch view for a layout, defaulting to the boundary layout.

get_packed_seq_params(

layout: core.context_parallel.layout.CPLayout | None = None,

) → megatron.core.packed_seq_params.PackedSeqParams | None

Return packed metadata for a layout, defaulting to the boundary layout.

core.context_parallel.utils._get_batch_on_this_cp_rank_contiguous(

batch: Dict[str, Any],

cp_group: torch.distributed.ProcessGroup,

) → Dict[str, Any]

Use contiguous CP shards while keeping dense attention-mask queries zigzag.

core.context_parallel.utils._get_batch_on_this_cp_rank_padded_zigzag(

batch: Dict[str, Any],

cp_group: torch.distributed.ProcessGroup,

rank_order_indices: torch.Tensor,

target_cu_seqlens_padded: torch.Tensor,

) → Dict[str, Any]

Build a packed zigzag shard padded for TE context-parallel attention.

core.context_parallel.utils._build_packed_seq_params(

batch: Dict[str, Any],

layout: core.context_parallel.layout.CPLayout,

cp_size: int,

tokens_per_sample: int | None,

use_logical_qkv_seqlens: bool = False,

pad_between_seqs: bool | None = None,

) → megatron.core.packed_seq_params.PackedSeqParams | None

Build packed sequence metadata for one physical CP layout.

core.context_parallel.utils.get_batches_on_this_cp_rank(

batch: Dict[str, Any],

boundary_layout: core.context_parallel.layout.CPLayout,

is_hybrid_cp: bool,

cp_group: torch.distributed.ProcessGroup,

additional_layouts: Iterable[core.context_parallel.layout.CPLayout] = (),

hybrid_cp_group_func: Callable[[int], torch.distributed.ProcessGroup] | None = None,

use_per_sequence_balancing: bool = False,

sequence_parallel: bool = False,

tp_group: torch.distributed.ProcessGroup | None = None,

tp_cp_group: torch.distributed.ProcessGroup | None = None,

tokens_per_sample: int | None = None,

) → core.context_parallel.utils.ContextParallelBatch

Partition a batch and prepare metadata for the requested CP layouts.

The input is already broadcast over TP but has not yet been partitioned over CP. A caller may request additional physical views when different consumers need different layouts. The boundary view is always included and is the default returned by ContextParallelBatch.

Packed non-hybrid CP batches need special handling when layout conversion requires a padded zigzag view or the standard sharder would balance the flattened sample instead of each sequence. The same rank ordering is used to build the zigzag batch tensors and their PackedSeqParams. When both layouts are requested, it also defines the activation-conversion plan. All other cases use the standard batch sharder.