`core.ssm.context_parallel.gdp_common`#

Shared execution path for Gated Delta Product context parallelism.

Module Contents#

Classes#

`GDPInputs`	Rank-local GDP operands passed to a chunkwise-CP backend.
`GDPChunkwiseContextParallel`	Connect a GDP backend’s saved context to PyTorch autograd.
`GDPChunkwiseCPBackend`	GDP local-kernel and autograd-adapter contract.

Functions#

gdp_chunkwise_context_parallel

Run GDP chunkwise CP through the selected local-kernel backend.

Data#

GDPInputGradients

API#

class core.ssm.context_parallel.gdp_common.GDPInputs#

Rank-local GDP operands passed to a chunkwise-CP backend.

q: torch.Tensor#: None

k: torch.Tensor#: None

v: torch.Tensor#: None

g: torch.Tensor#: None

beta: torch.Tensor#: None

cu_seqlens: torch.Tensor | None#: None

num_householder: int#: None

scale: float#: None

core.ssm.context_parallel.gdp_common.GDPInputGradients#: None

class core.ssm.context_parallel.gdp_common.GDPChunkwiseContextParallel#

Bases: torch.autograd.Function

Connect a GDP backend’s saved context to PyTorch autograd.

static forward( ctx, q, k, v, g, beta, cu_seqlens, num_householder, scale, cp_group, backend, preceding_rank_start, following_rank_stop, )#: Run chunkwise-CP forward and save the backend context for backward.

static backward(ctx, output_grad)#: Run chunkwise-CP backward using the saved backend context.

class core.ssm.context_parallel.gdp_common.GDPChunkwiseCPBackend#

Bases: megatron.core.ssm.context_parallel.chunkwise.LinearAttentionCPBackend[core.ssm.context_parallel.gdp_common.GDPInputs, megatron.core.ssm.context_parallel.chunkwise.LocalContextT, megatron.core.ssm.context_parallel.chunkwise.BackwardContextT, core.ssm.context_parallel.gdp_common.GDPInputGradients], typing.Protocol[megatron.core.ssm.context_parallel.chunkwise.LocalContextT, megatron.core.ssm.context_parallel.chunkwise.BackwardContextT]

GDP local-kernel and autograd-adapter contract.

autograd_function: type[core.ssm.context_parallel.gdp_common.GDPChunkwiseContextParallel]#: None

core.ssm.context_parallel.gdp_common.gdp_chunkwise_context_parallel( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor | None, num_householder: int, scale: float, cp_group: torch.distributed.ProcessGroup, backend: core.ssm.context_parallel.gdp_common.GDPChunkwiseCPBackend[megatron.core.ssm.context_parallel.chunkwise.LocalContextT, megatron.core.ssm.context_parallel.chunkwise.BackwardContextT], preceding_rank_start: int = 0, following_rank_stop: int | None = None, ) → torch.Tensor#: Run GDP chunkwise CP through the selected local-kernel backend.

core.ssm.context_parallel.gdp_common#