Module tp_bind

Module tp_bind 

Source
Expand description

2026-09-25: Applies a super::tp::KdaTpPlan: slices each KDA tensor to this rank before the binder sees it.

Owner: model-arch (GLM-5.3-Flash KDA). Invariants:

  • KdaShardedSource presents each tensor whose name maps to a plan entry (plan_name) as this rank’s bytes with a local shape, and passes every other name through unchanged.
  • shard_bytes refuses a tensor whose size disagrees with the plan’s full shape.

The slicing happens upstream of super::binding::bind_kda_weights, which validates the local shapes against the local config, so TP=1 and TP=2 run the same binder code.

Row slicing (HeadRows, ChannelRows) is one contiguous byte range. Column slicing (ChannelCols, only o_proj) is a strided gather: every row keeps its own middle slice. A row slice of o_proj has the same byte count at tp_size = 2 and the wrong contents, which is why the tests compare values.

Structs§

KdaShardedSource
2026-09-25: A KdaTensorSource that yields this rank’s slice of every planned KDA tensor. The loader wraps each layer’s source in it and passes it to bind_kda_weights.

Functions§

shard_bytes
2026-09-25: This rank’s bytes for one tensor, given its plan entry. full is the whole on-disk tensor; a size mismatch is an error. Returns a new buffer because a column slice is not contiguous.