`bridge.models.kimi.kimi_k3_pipeline`#

Stage-boundary packing for pipeline-parallel Kimi K3.

Module Contents#

Functions#

`bank_num_rows`	Return the number of AttnRes snapshots present before `layer_idx`.
`pack_stage_boundary`	Pack the AttnRes prefix and snapshot bank into one pipeline tensor.
`unpack_stage_boundary`	Unpack the prefix and snapshot bank received from a pipeline stage.

API#

bridge.models.kimi.kimi_k3_pipeline.bank_num_rows(layer_idx: int, block_size: int) → int#: Return the number of AttnRes snapshots present before layer_idx.

bridge.models.kimi.kimi_k3_pipeline.pack_stage_boundary( prefix_sum: torch.Tensor, block_residual: torch.Tensor, ) → torch.Tensor#: Pack the AttnRes prefix and snapshot bank into one pipeline tensor.

bridge.models.kimi.kimi_k3_pipeline.unpack_stage_boundary( packed: torch.Tensor, hidden_size: int, num_rows: int, ) → tuple[torch.Tensor, torch.Tensor]#: Unpack the prefix and snapshot bank received from a pipeline stage.