`core.ssm.context_parallel.chunkwise`#

Common chunkwise context-parallel interfaces for linear attention.

Module Contents#

Classes#

`CPForwardUnpackedSummary`	One shard’s unpacked `S_out = M S_in + delta_S` summary.
`CPForwardPackedSummary`	One shard’s packed `S_out = M S_in + delta_S` summary.
`CPBackwardUnpackedSummary`	One shard’s unpacked `dS_in = M^T dS_out + gamma` summary.
`CPBackwardPackedSummary`	One shard’s packed `dS_in = M^T dS_out + gamma` summary.
`CPSavedContext`	Backend state saved from forward for backward.
`PackedSequenceCPMetadata`	Rank-local metadata for chunkwise CP over packed sequences.
`LinearAttentionCPBackend`	Local kernel interface for one contiguous chunkwise-CP shard.
`CPForwardResult`	Results retained by a variant-specific autograd adapter after CP forward.

Functions#

`chunkwise_cp_forward`	Gather local summaries and apply the selected causal prefix.
`chunkwise_cp_backward`	Gather local adjoint summaries and apply the selected reverse-causal suffix.
`_all_gather_forward_summary`
`_all_gather_backward_summary`
`_slice_forward_summary`
`_slice_backward_summary`
`build_packed_sequence_cp_metadata`	Build rank-local packed-sequence metadata for chunkwise CP.
`_all_gather_unpacked_summary`	All-gather an unpacked state summary without packing its tensors.
`_all_gather_tensor`
`_validate_summary_tensor`	Validate a tensor communicated as part of a CP summary.

Data#

`CPForwardSummary`
`CPBackwardSummary`
`CPInputT`
`LocalContextT`
`BackwardContextT`
`CPGradientsT`

API#

class core.ssm.context_parallel.chunkwise.CPForwardUnpackedSummary#

One shard’s unpacked S_out = M S_in + delta_S summary.

transition: torch.Tensor#: None

state_update: torch.Tensor#: None

class core.ssm.context_parallel.chunkwise.CPForwardPackedSummary#

One shard’s packed S_out = M S_in + delta_S summary.

packed: torch.Tensor#: None

core.ssm.context_parallel.chunkwise.CPForwardSummary#: None

class core.ssm.context_parallel.chunkwise.CPBackwardUnpackedSummary#

One shard’s unpacked dS_in = M^T dS_out + gamma summary.

transition: torch.Tensor#: None

local_state_grad: torch.Tensor#: None

class core.ssm.context_parallel.chunkwise.CPBackwardPackedSummary#

One shard’s packed dS_in = M^T dS_out + gamma summary.

packed: torch.Tensor#: None

core.ssm.context_parallel.chunkwise.CPBackwardSummary#: None

class core.ssm.context_parallel.chunkwise.CPSavedContext#

Backend state saved from forward for backward.

metadata must not contain tensors; tensor state belongs in tensors so that the variant’s autograd adapter can save it with ctx.save_for_backward.

tensors: tuple[torch.Tensor, ...]#: None

metadata: object | None#: None

class core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata#

Rank-local metadata for chunkwise CP over packed sequences.

Parameters:

local_seq_idx – Per-token sequence IDs for the local contiguous shard.
local_cu_seqlens – Sequence boundaries within the local contiguous shard.
preceding_rank_start – Inclusive start of the forward-summary prefix.
following_rank_stop – Exclusive end of the backward-summary suffix.

local_seq_idx: torch.Tensor#: None

local_cu_seqlens: torch.Tensor#: None

preceding_rank_start: int#: None

following_rank_stop: int#: None

core.ssm.context_parallel.chunkwise.CPInputT#: ‘TypeVar(…)’

core.ssm.context_parallel.chunkwise.LocalContextT#: ‘TypeVar(…)’

core.ssm.context_parallel.chunkwise.BackwardContextT#: ‘TypeVar(…)’

core.ssm.context_parallel.chunkwise.CPGradientsT#: ‘TypeVar(…)’

class core.ssm.context_parallel.chunkwise.LinearAttentionCPBackend#

Bases: typing.Protocol[core.ssm.context_parallel.chunkwise.CPInputT, core.ssm.context_parallel.chunkwise.LocalContextT, core.ssm.context_parallel.chunkwise.BackwardContextT, core.ssm.context_parallel.chunkwise.CPGradientsT]

Local kernel interface for one contiguous chunkwise-CP shard.

Megatron owns all process-group operations and passes only selected summary slices to these methods.

cp_forward_prepare( inputs: core.ssm.context_parallel.chunkwise.CPInputT, ) → tuple[core.ssm.context_parallel.chunkwise.CPForwardSummary, core.ssm.context_parallel.chunkwise.LocalContextT]#: Compute the local state summary with zero incoming state.

cp_forward_apply( local_context: core.ssm.context_parallel.chunkwise.LocalContextT, preceding_summaries: core.ssm.context_parallel.chunkwise.CPForwardSummary, ) → tuple[torch.Tensor, core.ssm.context_parallel.chunkwise.CPSavedContext]#: Compose the preceding summaries and compute the local output.

cp_backward_prepare( output_grad: torch.Tensor, saved_context: core.ssm.context_parallel.chunkwise.CPSavedContext, ) → tuple[core.ssm.context_parallel.chunkwise.CPBackwardSummary, core.ssm.context_parallel.chunkwise.BackwardContextT]#: Compute local gamma with zero outgoing-state gradient.

cp_backward_apply( backward_context: core.ssm.context_parallel.chunkwise.BackwardContextT, following_summaries: core.ssm.context_parallel.chunkwise.CPBackwardSummary, ) → core.ssm.context_parallel.chunkwise.CPGradientsT#: Compose the reverse-causal suffix and finish the local input gradients.

class core.ssm.context_parallel.chunkwise.CPForwardResult#

Results retained by a variant-specific autograd adapter after CP forward.

output: torch.Tensor#: None

saved_context: core.ssm.context_parallel.chunkwise.CPSavedContext#: None

core.ssm.context_parallel.chunkwise.chunkwise_cp_forward( backend: core.ssm.context_parallel.chunkwise.LinearAttentionCPBackend[core.ssm.context_parallel.chunkwise.CPInputT, core.ssm.context_parallel.chunkwise.LocalContextT, core.ssm.context_parallel.chunkwise.BackwardContextT, core.ssm.context_parallel.chunkwise.CPGradientsT], inputs: core.ssm.context_parallel.chunkwise.CPInputT, cp_group: torch.distributed.ProcessGroup, preceding_slice: slice, ) → core.ssm.context_parallel.chunkwise.CPForwardResult#: Gather local summaries and apply the selected causal prefix.

core.ssm.context_parallel.chunkwise.chunkwise_cp_backward( backend: core.ssm.context_parallel.chunkwise.LinearAttentionCPBackend[core.ssm.context_parallel.chunkwise.CPInputT, core.ssm.context_parallel.chunkwise.LocalContextT, core.ssm.context_parallel.chunkwise.BackwardContextT, core.ssm.context_parallel.chunkwise.CPGradientsT], output_grad: torch.Tensor, saved_context: core.ssm.context_parallel.chunkwise.CPSavedContext, cp_group: torch.distributed.ProcessGroup, following_slice: slice, ) → core.ssm.context_parallel.chunkwise.CPGradientsT#: Gather local adjoint summaries and apply the selected reverse-causal suffix.

core.ssm.context_parallel.chunkwise._all_gather_forward_summary( local_summary: core.ssm.context_parallel.chunkwise.CPForwardSummary, cp_group: torch.distributed.ProcessGroup, ) → core.ssm.context_parallel.chunkwise.CPForwardSummary#

core.ssm.context_parallel.chunkwise._all_gather_backward_summary( local_summary: core.ssm.context_parallel.chunkwise.CPBackwardSummary, cp_group: torch.distributed.ProcessGroup, ) → core.ssm.context_parallel.chunkwise.CPBackwardSummary#

core.ssm.context_parallel.chunkwise._slice_forward_summary( summary: core.ssm.context_parallel.chunkwise.CPForwardSummary, rank_slice: slice, ) → core.ssm.context_parallel.chunkwise.CPForwardSummary#

core.ssm.context_parallel.chunkwise._slice_backward_summary( summary: core.ssm.context_parallel.chunkwise.CPBackwardSummary, rank_slice: slice, ) → core.ssm.context_parallel.chunkwise.CPBackwardSummary#

core.ssm.context_parallel.chunkwise.build_packed_sequence_cp_metadata( global_seq_idx: torch.Tensor, cp_rank: int, cp_size: int, ) → core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata#

Build rank-local packed-sequence metadata for chunkwise CP.