Module tp_bind 

Applying the TP plan: the shard copies, as an upstream adapter so the binder is untouched. Applying super::tp::KdaTpPlan — the shard COPIES, at last.

Why this is an adapter and not a change to the binder

super::binding::bind_kda_weights is proven: exact tensor set, exact dtypes, exact shapes, gated numerically on real weights. Teaching it about TP would mean rewriting its validation to check on-disk (full) shapes against a per-rank config — i.e. editing the one piece of this lane that has never been wrong.

So the slicing happens upstream. KdaShardedSource wraps any KdaTensorSource and hands the binder bytes that are already this rank’s, with local shapes. The binder then validates them against the local config exactly as it always has and cannot tell the difference — which is the point: TP=1 and TP=2 run the same binder code path, so there is no “works at TP=1, silently differs at TP=2” seam for a bug to live in.

🪤 Two slicing shapes, and the wrong one is not a shape error

• Row slicing (HeadRows / ChannelRows) is a contiguous byte range: this rank’s rows sit next to each other on disk.
• Column slicing (ChannelCols, i.e. o_proj alone) is a STRIDED gather — every row contributes its own middle slice. Row-slicing o_proj instead yields a well-formed [hidden/tp, heads*head_dim] tensor of real numbers and a wrong output.

Both produce the same LOCAL element count at tp_size = 2 when the tensor is square-ish, so a size check cannot separate them. The test below separates them by value.

Structs

KdaShardedSource

A KdaTensorSource that yields this rank’s slice of every KDA tensor.

Functions

shard_bytes

This rank’s bytes for one tensor, given its plan entry.