Module tp_shard 

Tensor-parallel weight sharding helpers.

Megatron-style TP slices each weight tensor along one of two axes:

• Column-parallel (Q/K/V proj, gate_proj, up_proj, lm_head): weight shape [out, in] becomes [out / tp, in]. Rank r keeps rows [r * out / tp, (r + 1) * out / tp). This is a single contiguous slice in row-major layout — 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). Per-row strided copy because the surviving slice is non-contiguous in row-major layout.

1D per-output vectors (q_norm_full, k_norm_full, gate_proj bias, etc.) shard with the same axis as their associated GEMM’s column-parallel output.

All shard helpers operate on BF16 weights before NVFP4 quantization; sharding the packed FP4 storage + FP8 scales is mechanical but adds two more axes to bookkeep, and pre-quant slicing keeps the existing quantize path untouched.

Structs

TpAttentionDims

Pre-TP-shard attention dimensions reconstructed from config.

TpGdnDims

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

TpMoeDims

Pre-TP-shard dimensions for MoE expert projections. Unlike attention, main.rs does NOT divide moe_intermediate_size by tp_size, so full_inter == config.moe_intermediate_size. Local size is computed here for downstream callers.

Enums

TpShardKind

TP shard kind for a 2D BF16 weight [out_dim, in_dim].

Functions

load_qk_norms_tp

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

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

Shard a 1D BF16 vector [dim] (e.g. q_norm_full, gate_proj bias) on dim 0. Used for per-output vectors that pair with column-parallel GEMMs.

shard_dense_bf16

Shard a BF16 dense weight [out_dim, in_dim] according to kind.

shard_dense_weight

Convenience wrapper: shard a DenseWeight BF16 tensor. The source weight is freed by the caller — this fn allocates a new device buffer.

shard_fp8_block_scaled

Shard an FP8 block-scaled weight. weight is [N, K] FP8 bytes; row_scale is [N/block_size, K/block_size] FP32 (widened at load). Both slice on the same axis at block granularity.

shard_gdn_ba_rows

Shard the BA gate BF16 weight [2*full_nv, h] to [2*local_nv, h].

shard_gdn_conv_rows

Shard the depthwise conv1d BF16 weight [full_conv_dim, d_conv] to [local_conv_dim, d_conv]. Channels ARE the QKV channels (one filter per channel), so this uses the SAME [Q|K|V] segment pattern as the QKV in-projection — the conv is NOT replicated across ranks.

shard_gdn_out_proj_row_parallel

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 partial products are summed with an all-reduce after the GEMM (mirrors attention o_proj). Returns (ptr, h, local_value_dim).

shard_gdn_qkv_rows

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

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

Shard a per-value-head 1D vector on the value-head axis. Handles BF16 (norm, [full_nv*vd] → [local_nv*vd] with elem_bytes = 2, unit = vd) and FP32 (a_log / dt_bias, [full_nv] → [local_nv] with elem_bytes = 4, unit = 1). unit is the number of elements per value head. Returns (ptr, local_len_elems).

shard_quantized_nvfp4

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).