core.packed_seq_params

Module Contents

Classes

PackedSeqParams

Packed sequence boundaries and MCore’s per-microbatch parallel topology.

Functions

resolve_cp_group	Resolve and validate the context-parallel group for this microbatch.
resolve_tp_cp_group	Resolve the TP x CP group matching this microbatch’s context-parallel group.

API

class core.packed_seq_params.PackedSeqParams

Packed sequence boundaries and MCore’s per-microbatch parallel topology.

Attention and fused RoPE kernels consume the sequence metadata. cp_group, local_cp_size and tp_cp_group are MCore-only fields and must be filtered before forwarding arguments to Transformer Engine. Runtime CP includes a real singleton group; local_cp_size requires a matching cp_group.

qkv_format: str

None

cu_seqlens_q: torch.Tensor

None

cu_seqlens_kv: torch.Tensor

None

cu_seqlens_q_padded: torch.Tensor

None

cu_seqlens_kv_padded: torch.Tensor

None

max_seqlen_q: int

None

max_seqlen_kv: int

None

local_cp_size: int | None

None

cp_group: torch.distributed.ProcessGroup | None

None

tp_cp_group: torch.distributed.ProcessGroup | None

None

total_tokens: int | None

None

seq_idx: torch.Tensor

None

tokens_per_sample: int

None

pad_between_seqs: bool

None

cp_scatter_cache: object

None

__post_init__()

Pre-compute seq_idx for Mamba mixer CUDA graph compatibility.

If total_tokens is 16 (for example), this method takes packed_seq_params.cu_seqlens_q_padded (or cu_seqlens_q) which is of the form [0, 5, 7, 11] and returns a tensor of the form [0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3], which is [0](5-0) + [1](7-5) + [2](11-7) + [3](16-11) In the above example, there are three sequences in the pack. In general, the output has an additional sequence index (e.g. 0, 1, 2, 3) so that any tokens beyond the last padded input sequence are accounted for as an extra sequence. However, If cu_seqlens_q_padded[-1] == max_seqlen then this additional sequence index will not be included.

core.packed_seq_params.resolve_cp_group(

static_cp_group: torch.distributed.ProcessGroup | None,

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

) → torch.distributed.ProcessGroup | None

Resolve and validate the context-parallel group for this microbatch.

core.packed_seq_params.resolve_tp_cp_group(

static_tp_cp_group: torch.distributed.ProcessGroup | None,

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

*,

sequence_parallel: bool,

tp_group: torch.distributed.ProcessGroup,

) → torch.distributed.ProcessGroup | None

Resolve the TP x CP group matching this microbatch’s context-parallel group.