core.context_parallel.sequence_roll

Contiguous packed sequence rolls with explicit groups and differentiable boundary exchange.

Module Contents

Classes

_RollLeft

One differentiable left roll, including the reverse neighbor exchange.

Functions

_continuation_mask	Return the local rows whose next token belongs to the same real document.
roll_contiguous	Roll a dense or packed contiguous shard by one, preserving real document boundaries.
roll_contiguous_fields	Roll non-differentiable token fields together, preserving each dtype and fill value.

API

core.context_parallel.sequence_roll._continuation_mask(tensor, packed_seq_params, cp_group, dims)

Return the local rows whose next token belongs to the same real document.

class core.context_parallel.sequence_roll._RollLeft

Bases: torch.autograd.Function

One differentiable left roll, including the reverse neighbor exchange.

static forward(ctx, tensor, group, dims)

Shift local values and receive the next rank’s first row.

static backward(ctx, grad)

Return the boundary gradient to its original token owner.

core.context_parallel.sequence_roll.roll_contiguous(

tensor,

dims,

cp_group,

packed_seq_params=None,

fill_value=0,

)

Roll a dense or packed contiguous shard by one, preserving real document boundaries.

core.context_parallel.sequence_roll.roll_contiguous_fields(

tensors,

cp_group,

packed_seq_params=None,

fill_values=None,

)

Roll non-differentiable token fields together, preserving each dtype and fill value.

This batches one depth’s IDs/positions/masks; activations continue to use the differentiable roll above. No cross-layer cache or model-wide state is introduced.