shard_quantized_nvfp4

Function shard_quantized_nvfp4 

Source
pub fn shard_quantized_nvfp4(
    src: &QuantizedWeight,
    out_dim: usize,
    in_dim: usize,
    kind: TpShardKind,
    tp_rank: usize,
    tp_size: usize,
    group_size: usize,
    gpu: &dyn GpuBackend,
) -> Result<QuantizedWeight>
Expand description

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

out_dim (= N) and in_dim (= K) are pre-shard, full-tensor dims. Returns a freshly-allocated, sharded QuantizedWeight. Caller frees the source if tp_size > 1.