Module tp_shard 

2026-09-25: Tensor-parallel weight sharding helpers.

Each weight tensor is sliced along one of two axes:

• Column-parallel (Q/K/V proj, gate_proj, up_proj): weight shape [out, in] becomes [out / tp, in]. Rank r keeps rows [r * out / tp, (r + 1) * out / tp), one contiguous slice in row-major layout, so one copy_d2d.

• Row-parallel (O proj, down_proj): weight [out, in] becomes [out, in / tp]. Rank r keeps cols [r * in / tp, (r + 1) * in / tp), one strided copy per row.

1D per-output vectors shard along the output axis of their column-parallel GEMM (shard_dense_1d_bf16).

The BF16 primitives are here; quant_shard.rs slices weights that are already NVFP4 or FP8 block-scaled, and gdn.rs the segmented Gated-DeltaNet projections.

Owner: model-arch (tensor parallelism). Invariants:

• With tp_size <= 1 every shard function returns its source unchanged and allocates nothing.

Structs

TpAttentionDims

2026-09-25: Pre-TP-shard attention dimensions reconstructed from config.

TpGdnDims

2026-09-25: Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from config.

TpMoeDims

2026-09-25: Pre-TP-shard dimensions for MoE expert projections. Unlike the head counts, serve_phases/topology.rs does not divide moe_intermediate_size, so full_inter == config.moe_intermediate_size and the local size is computed here.

Enums

TpShardKind

2026-09-25: TP shard kind for a 2D weight [out_dim, in_dim].

Functions

load_qk_norms_tp

2026-09-25: Q/K-norm 1D shard pair. The closure receives (name, full_dim) and returns the loader’s sharded norm — typically a DenseWeight. q_norm is sharded against full_q_n; k_norm against full_kv_n. Returns (q_norm, k_norm).

load_qkvo_tp

2026-09-25: Sequence the four Q/K/V/O loads via a loader-supplied closure. The closure receives (name, full_out, full_in, kind) and returns the loader’s representation of that projection (BF16 dense, NVFP4 quantized, FP8 block-scaled — varies by format).

shard_dense_1d_bf16

2026-09-25: Shard a 1D BF16 vector [dim] on dim 0, for per-output vectors of column-parallel GEMMs. Same ownership rule as shard_dense_bf16.

shard_dense_bf16

2026-09-25: Shard a BF16 dense weight [out_dim, in_dim] according to kind.

shard_dense_weight

2026-09-25: shard_dense_bf16 on a DenseWeight, with the same ownership rule: a new buffer only when a shard happened, and then the caller frees the source.

shard_fp8_block_scaled

2026-09-25: Shard an FP8 block-scaled weight. weight is [N, K] FP8 bytes; row_scale is [N/block_size, K/block_size] FP32. Both slice on the same axis at block granularity. Same ownership rule as shard_quantized_nvfp4.

shard_gdn_ba_rows

2026-09-25: Shard the BA gate BF16 weight [2*full_nv, h] to [2*local_nv, h].

shard_gdn_conv_rows

2026-09-25: Shard the depthwise conv1d BF16 weight [full_conv_dim, d_conv] to [local_conv_dim, d_conv]. Its channels are the QKV channels (one filter per channel), so it uses the same [Q|K|V] segments as the QKV in-projection.

shard_gdn_out_proj_row_parallel

2026-09-25: Shard the out_proj BF16 weight [h, full_value_dim] row-parallel on its input dim (value_dim). Rank r keeps columns [r*local_value_dim, (r+1)*local_value_dim) of every output row; the layer sums the partial products with an all-reduce after the GEMM. Returns (ptr, h, local_value_dim).

shard_gdn_qkv_rows

2026-09-25: Shard the [Q|K|V] (in_proj_qkv) BF16 weight [full_conv_dim, h] to the local rank’s [local_conv_dim, h], slicing Q, K and V independently by the local head range. Returns (ptr, local_rows, h).

shard_gdn_qkvz_rows

2026-09-25: Shard the concatenated [Q|K|V|Z] (in_proj_qkvz) BF16 weight [full_qkvz_out, h] to the local rank’s [local_qkvz_out, h], slicing all four segments independently. Returns (ptr, local_rows, h).

shard_gdn_value_vector

2026-09-25: Shard a per-value-head 1D vector on the value-head axis: [full_nv * unit] → [local_nv * unit], where unit is the number of elements per value head and elem_bytes their size. The loaders call it for a_log / dt_bias (unit = 1, elem_bytes = 4). Returns (ptr, local_len_elems).

shard_quantized_nvfp4

2026-09-25: Shard an NVFP4-quantized weight. The packed weight is [N, K/2] u8; the per-group scale is [N, K/group_size] u8 (FP8); weight_scale_2 is a per-tensor f32 (replicated across all ranks).