metrale_model_arch/tp_shard/
quant_shard.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Shard primitives for weights that are already quantized:
4//! NVFP4 (`shard_quantized_nvfp4`) and FP8 block-scaled
5//! (`shard_fp8_block_scaled`).
6//!
7//! Owner: model-arch (tensor parallelism).
8//! Invariants:
9//! - With `tp_size <= 1` or `Replicated`, both return a copy of the source
10//!   (same device pointers) and allocate nothing.
11
12use anyhow::{Result, ensure};
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::TpShardKind;
16use metrale_model_layers::weight_map::{Fp8Weight, QuantizedWeight};
17
18// 2026-09-25: The layouts these functions assume:
19//
20//   - NVFP4 packed: `weight` is `[N, K/2]` u8, `weight_scale` is
21//     `[N, K/group_size]` u8 (FP8 E4M3), `weight_scale_2` is a single
22//     f32. Both byte tensors slice on the same N (column-parallel) or
23//     K axis (row-parallel); the per-tensor scale is replicated.
24//
25//   - FP8 block-scaled: `weight` is `[N, K]` FP8 bytes, `row_scale` is
26//     `[N/block_size, K/block_size]` FP32. Both slice on N
27//     (column-parallel) or K (row-parallel) at block granularity.
28
29/// 2026-09-25: Shard an NVFP4-quantized weight. The packed weight is `[N, K/2]` u8;
30/// the per-group scale is `[N, K/group_size]` u8 (FP8); `weight_scale_2`
31/// is a per-tensor f32 (replicated across all ranks).
32///
33/// `out_dim` (= N) and `in_dim` (= K) are pre-shard, full-tensor dims. With
34/// `tp_size <= 1` or `Replicated` it returns a copy of `src`; otherwise newly
35/// allocated buffers, and the caller frees the source.
36pub fn shard_quantized_nvfp4(
37    src: &QuantizedWeight,
38    out_dim: usize,
39    in_dim: usize,
40    kind: TpShardKind,
41    tp_rank: usize,
42    tp_size: usize,
43    group_size: usize,
44    gpu: &dyn GpuBackend,
45) -> Result<QuantizedWeight> {
46    if tp_size <= 1 || kind == TpShardKind::Replicated {
47        return Ok(*src);
48    }
49    ensure!(
50        in_dim.is_multiple_of(group_size),
51        "NVFP4 in_dim {in_dim} not divisible by group_size {group_size}",
52    );
53    let half_in = in_dim / 2;
54    let scale_in = in_dim / group_size;
55    match kind {
56        TpShardKind::Replicated => unreachable!("handled above"),
57        TpShardKind::ColumnParallel => {
58            ensure!(
59                out_dim.is_multiple_of(tp_size),
60                "NVFP4 ColumnParallel: out_dim {out_dim} not divisible by tp_size {tp_size}",
61            );
62            let local_out = out_dim / tp_size;
63            // 2026-09-25: Packed weight: [N, K/2] u8 → this rank's local_out rows.
64            let w_row_bytes = half_in;
65            let w_local_bytes = local_out * w_row_bytes;
66            let w_dst = gpu.alloc(w_local_bytes)?;
67            let w_offset = tp_rank * local_out * w_row_bytes;
68            gpu.copy_d2d(
69                DevicePtr(src.weight.0 + w_offset as u64),
70                w_dst,
71                w_local_bytes,
72            )?;
73            // 2026-09-25: Scale: [N, K/group_size] u8 → the same rows.
74            let s_row_bytes = scale_in;
75            let s_local_bytes = local_out * s_row_bytes;
76            let s_dst = gpu.alloc(s_local_bytes)?;
77            let s_offset = tp_rank * local_out * s_row_bytes;
78            gpu.copy_d2d(
79                DevicePtr(src.weight_scale.0 + s_offset as u64),
80                s_dst,
81                s_local_bytes,
82            )?;
83            Ok(QuantizedWeight {
84                weight: w_dst,
85                weight_scale: s_dst,
86                weight_scale_2: src.weight_scale_2,
87                input_scale: src.input_scale,
88                weight_scale_2_vec: src.weight_scale_2_vec,
89                act: src.act,
90            })
91        }
92        TpShardKind::RowParallel => {
93            ensure!(
94                in_dim.is_multiple_of(tp_size),
95                "NVFP4 RowParallel: in_dim {in_dim} not divisible by tp_size {tp_size}",
96            );
97            ensure!(
98                half_in.is_multiple_of(tp_size),
99                "NVFP4 RowParallel: half_in {half_in} (=K/2) not divisible by tp_size {tp_size}",
100            );
101            ensure!(
102                scale_in.is_multiple_of(tp_size),
103                "NVFP4 RowParallel: scale_in {scale_in} (=K/group_size) not divisible by tp_size {tp_size}",
104            );
105            let local_w_in = half_in / tp_size;
106            let local_s_in = scale_in / tp_size;
107            // 2026-09-25: Per-row strided copy on the K axis for both weight and scale.
108            let w_local_row_bytes = local_w_in;
109            let w_src_row_bytes = half_in;
110            let w_local_bytes = out_dim * w_local_row_bytes;
111            let w_dst = gpu.alloc(w_local_bytes)?;
112            let w_col_offset = tp_rank * w_local_row_bytes;
113            for r in 0..out_dim {
114                gpu.copy_d2d(
115                    DevicePtr(src.weight.0 + (r * w_src_row_bytes + w_col_offset) as u64),
116                    DevicePtr(w_dst.0 + (r * w_local_row_bytes) as u64),
117                    w_local_row_bytes,
118                )?;
119            }
120            let s_local_row_bytes = local_s_in;
121            let s_src_row_bytes = scale_in;
122            let s_local_bytes = out_dim * s_local_row_bytes;
123            let s_dst = gpu.alloc(s_local_bytes)?;
124            let s_col_offset = tp_rank * s_local_row_bytes;
125            for r in 0..out_dim {
126                gpu.copy_d2d(
127                    DevicePtr(src.weight_scale.0 + (r * s_src_row_bytes + s_col_offset) as u64),
128                    DevicePtr(s_dst.0 + (r * s_local_row_bytes) as u64),
129                    s_local_row_bytes,
130                )?;
131            }
132            Ok(QuantizedWeight {
133                weight: w_dst,
134                weight_scale: s_dst,
135                weight_scale_2: src.weight_scale_2,
136                input_scale: src.input_scale,
137                weight_scale_2_vec: src.weight_scale_2_vec,
138                act: src.act,
139            })
140        }
141    }
142}
143
144/// 2026-09-25: Shard an FP8 block-scaled weight. `weight` is `[N, K]` FP8 bytes;
145/// `row_scale` is `[N/block_size, K/block_size]` FP32. Both slice on the same
146/// axis at block granularity. Same ownership rule as `shard_quantized_nvfp4`.
147pub fn shard_fp8_block_scaled(
148    src: &Fp8Weight,
149    kind: TpShardKind,
150    tp_rank: usize,
151    tp_size: usize,
152    block_size: usize,
153    gpu: &dyn GpuBackend,
154) -> Result<Fp8Weight> {
155    if tp_size <= 1 || kind == TpShardKind::Replicated {
156        return Ok(*src);
157    }
158    let n = src.n as usize;
159    let k = src.k as usize;
160    ensure!(
161        n.is_multiple_of(block_size),
162        "FP8 N {n} not divisible by block_size {block_size}",
163    );
164    ensure!(
165        k.is_multiple_of(block_size),
166        "FP8 K {k} not divisible by block_size {block_size}",
167    );
168    let scale_n = n / block_size;
169    let scale_k = k / block_size;
170    let scale_elem_bytes = 4usize;
171    match kind {
172        TpShardKind::Replicated => unreachable!("handled above"),
173        TpShardKind::ColumnParallel => {
174            ensure!(
175                n.is_multiple_of(tp_size),
176                "FP8 ColumnParallel: N {n} not divisible by tp_size {tp_size}",
177            );
178            ensure!(
179                scale_n.is_multiple_of(tp_size),
180                "FP8 ColumnParallel: scale_n {scale_n} not divisible by tp_size {tp_size}",
181            );
182            let local_n = n / tp_size;
183            let local_scale_n = scale_n / tp_size;
184            // 2026-09-25: Weight: [N, K] u8 → this rank's local_n rows.
185            let w_row_bytes = k;
186            let w_local_bytes = local_n * w_row_bytes;
187            let w_dst = gpu.alloc(w_local_bytes)?;
188            let w_offset = tp_rank * local_n * w_row_bytes;
189            gpu.copy_d2d(
190                DevicePtr(src.weight.0 + w_offset as u64),
191                w_dst,
192                w_local_bytes,
193            )?;
194            // 2026-09-25: Scale: [scale_n, scale_k] FP32 → this rank's local_scale_n rows.
195            let s_row_bytes = scale_k * scale_elem_bytes;
196            let s_local_bytes = local_scale_n * s_row_bytes;
197            let s_dst = gpu.alloc(s_local_bytes)?;
198            let s_offset = tp_rank * local_scale_n * s_row_bytes;
199            gpu.copy_d2d(
200                DevicePtr(src.row_scale.0 + s_offset as u64),
201                s_dst,
202                s_local_bytes,
203            )?;
204            Ok(Fp8Weight {
205                weight: w_dst,
206                row_scale: s_dst,
207                n: local_n as u32,
208                k: src.k,
209                scale_format: src.scale_format,
210            })
211        }
212        TpShardKind::RowParallel => {
213            ensure!(
214                k.is_multiple_of(tp_size),
215                "FP8 RowParallel: K {k} not divisible by tp_size {tp_size}",
216            );
217            ensure!(
218                scale_k.is_multiple_of(tp_size),
219                "FP8 RowParallel: scale_k {scale_k} not divisible by tp_size {tp_size}",
220            );
221            let local_k = k / tp_size;
222            let local_scale_k = scale_k / tp_size;
223            let w_local_row_bytes = local_k;
224            let w_src_row_bytes = k;
225            let w_local_bytes = n * w_local_row_bytes;
226            let w_dst = gpu.alloc(w_local_bytes)?;
227            let w_col_offset = tp_rank * w_local_row_bytes;
228            for r in 0..n {
229                gpu.copy_d2d(
230                    DevicePtr(src.weight.0 + (r * w_src_row_bytes + w_col_offset) as u64),
231                    DevicePtr(w_dst.0 + (r * w_local_row_bytes) as u64),
232                    w_local_row_bytes,
233                )?;
234            }
235            let s_local_row_bytes = local_scale_k * scale_elem_bytes;
236            let s_src_row_bytes = scale_k * scale_elem_bytes;
237            let s_local_bytes = scale_n * s_local_row_bytes;
238            let s_dst = gpu.alloc(s_local_bytes)?;
239            let s_col_offset = tp_rank * s_local_row_bytes;
240            for r in 0..scale_n {
241                gpu.copy_d2d(
242                    DevicePtr(src.row_scale.0 + (r * s_src_row_bytes + s_col_offset) as u64),
243                    DevicePtr(s_dst.0 + (r * s_local_row_bytes) as u64),
244                    s_local_row_bytes,
245                )?;
246            }
247            Ok(Fp8Weight {
248                weight: w_dst,
249                row_scale: s_dst,
250                n: src.n,
251                k: local_k as u32,
252                scale_format: src.scale_format,
253            })
254        }
255    }
256}
metrale_model_arch/tp_shard/quant_shard.rs

metrale_model_arch/tp_shard/
quant_shard.rs
