spark_model/tp_shard/
quant_shard.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Format-specific shard primitives for pre-quantized weights.
4
5use anyhow::{Result, ensure};
6use spark_runtime::gpu::{DevicePtr, GpuBackend};
7
8use super::TpShardKind;
9use crate::weight_map::{Fp8Weight, QuantizedWeight};
10
11// ── Format-specific shard primitives ───────────────────────────────
12//
13// `shard_dense_bf16` (above) handles the BF16 case used by every
14// loader's "load BF16 → shard → quantize" path. The two below handle
15// the formats that ship pre-quantized on disk:
16//
17//   - NVFP4 packed: `weight` is `[N, K/2]` u8, `weight_scale` is
18//     `[N, K/group_size]` u8 (FP8 E4M3), `weight_scale_2` is a single
19//     f32. Both byte tensors slice on the same N (column-parallel) or
20//     K axis (row-parallel); the per-tensor scale is replicated.
21//
22//   - FP8 block-scaled: `weight` is `[N, K]` FP8 bytes, `row_scale` is
23//     `[N/block_size, K/block_size]` BF16. Both slice on N
24//     (column-parallel) or K (row-parallel) at block granularity.
25//
26// All helpers preserve the "tp_size == 1 returns source untouched" fast
27// path matching `shard_dense_bf16`.
28
29/// 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.
34/// Returns a freshly-allocated, sharded `QuantizedWeight`. Caller frees
35/// the source if `tp_size > 1`.
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; // packed weight stride
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            // Packed weight: [N, K/2] u8 → slice top 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            // Scale: [N, K/group_size] u8 → slice same N axis.
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            })
90        }
91        TpShardKind::RowParallel => {
92            ensure!(
93                in_dim.is_multiple_of(tp_size),
94                "NVFP4 RowParallel: in_dim {in_dim} not divisible by tp_size {tp_size}",
95            );
96            ensure!(
97                half_in.is_multiple_of(tp_size),
98                "NVFP4 RowParallel: half_in {half_in} (=K/2) not divisible by tp_size {tp_size}",
99            );
100            ensure!(
101                scale_in.is_multiple_of(tp_size),
102                "NVFP4 RowParallel: scale_in {scale_in} (=K/group_size) not divisible by tp_size {tp_size}",
103            );
104            let local_w_in = half_in / tp_size;
105            let local_s_in = scale_in / tp_size;
106            // Per-row strided copy on the K axis for both weight and scale.
107            let w_local_row_bytes = local_w_in;
108            let w_src_row_bytes = half_in;
109            let w_local_bytes = out_dim * w_local_row_bytes;
110            let w_dst = gpu.alloc(w_local_bytes)?;
111            let w_col_offset = tp_rank * w_local_row_bytes;
112            for r in 0..out_dim {
113                gpu.copy_d2d(
114                    DevicePtr(src.weight.0 + (r * w_src_row_bytes + w_col_offset) as u64),
115                    DevicePtr(w_dst.0 + (r * w_local_row_bytes) as u64),
116                    w_local_row_bytes,
117                )?;
118            }
119            let s_local_row_bytes = local_s_in;
120            let s_src_row_bytes = scale_in;
121            let s_local_bytes = out_dim * s_local_row_bytes;
122            let s_dst = gpu.alloc(s_local_bytes)?;
123            let s_col_offset = tp_rank * s_local_row_bytes;
124            for r in 0..out_dim {
125                gpu.copy_d2d(
126                    DevicePtr(src.weight_scale.0 + (r * s_src_row_bytes + s_col_offset) as u64),
127                    DevicePtr(s_dst.0 + (r * s_local_row_bytes) as u64),
128                    s_local_row_bytes,
129                )?;
130            }
131            Ok(QuantizedWeight {
132                weight: w_dst,
133                weight_scale: s_dst,
134                weight_scale_2: src.weight_scale_2,
135                input_scale: src.input_scale,
136                weight_scale_2_vec: src.weight_scale_2_vec,
137            })
138        }
139    }
140}
141
142/// Shard an FP8 block-scaled weight. `weight` is `[N, K]` FP8 bytes;
143/// `row_scale` is `[N/block_size, K/block_size]` FP32 (widened at load).
144/// Both slice on the same axis at block granularity.
145pub fn shard_fp8_block_scaled(
146    src: &Fp8Weight,
147    kind: TpShardKind,
148    tp_rank: usize,
149    tp_size: usize,
150    block_size: usize,
151    gpu: &dyn GpuBackend,
152) -> Result<Fp8Weight> {
153    if tp_size <= 1 || kind == TpShardKind::Replicated {
154        return Ok(*src);
155    }
156    let n = src.n as usize;
157    let k = src.k as usize;
158    ensure!(
159        n.is_multiple_of(block_size),
160        "FP8 N {n} not divisible by block_size {block_size}",
161    );
162    ensure!(
163        k.is_multiple_of(block_size),
164        "FP8 K {k} not divisible by block_size {block_size}",
165    );
166    let scale_n = n / block_size;
167    let scale_k = k / block_size;
168    // Block scale is FP32 (widened from the checkpoint at load); 4 bytes/elem.
169    let scale_elem_bytes = 4usize;
170    match kind {
171        TpShardKind::Replicated => unreachable!("handled above"),
172        TpShardKind::ColumnParallel => {
173            ensure!(
174                n.is_multiple_of(tp_size),
175                "FP8 ColumnParallel: N {n} not divisible by tp_size {tp_size}",
176            );
177            ensure!(
178                scale_n.is_multiple_of(tp_size),
179                "FP8 ColumnParallel: scale_n {scale_n} not divisible by tp_size {tp_size}",
180            );
181            let local_n = n / tp_size;
182            let local_scale_n = scale_n / tp_size;
183            // Weight: [N, K] u8 → slice top local_n rows.
184            let w_row_bytes = k;
185            let w_local_bytes = local_n * w_row_bytes;
186            let w_dst = gpu.alloc(w_local_bytes)?;
187            let w_offset = tp_rank * local_n * w_row_bytes;
188            gpu.copy_d2d(
189                DevicePtr(src.weight.0 + w_offset as u64),
190                w_dst,
191                w_local_bytes,
192            )?;
193            // Scale: [scale_n, scale_k] BF16 → slice top local_scale_n rows.
194            let s_row_bytes = scale_k * scale_elem_bytes;
195            let s_local_bytes = local_scale_n * s_row_bytes;
196            let s_dst = gpu.alloc(s_local_bytes)?;
197            let s_offset = tp_rank * local_scale_n * s_row_bytes;
198            gpu.copy_d2d(
199                DevicePtr(src.row_scale.0 + s_offset as u64),
200                s_dst,
201                s_local_bytes,
202            )?;
203            Ok(Fp8Weight {
204                weight: w_dst,
205                row_scale: s_dst,
206                n: local_n as u32,
207                k: src.k,
208                scale_format: src.scale_format,
209            })
210        }
211        TpShardKind::RowParallel => {
212            ensure!(
213                k.is_multiple_of(tp_size),
214                "FP8 RowParallel: K {k} not divisible by tp_size {tp_size}",
215            );
216            ensure!(
217                scale_k.is_multiple_of(tp_size),
218                "FP8 RowParallel: scale_k {scale_k} not divisible by tp_size {tp_size}",
219            );
220            let local_k = k / tp_size;
221            let local_scale_k = scale_k / tp_size;
222            // Weight: per-row strided on K.
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            // Scale: per-row strided on scale_k.
236            let s_local_row_bytes = local_scale_k * scale_elem_bytes;
237            let s_src_row_bytes = scale_k * scale_elem_bytes;
238            let s_local_bytes = scale_n * s_local_row_bytes;
239            let s_dst = gpu.alloc(s_local_bytes)?;
240            let s_col_offset = tp_rank * s_local_row_bytes;
241            for r in 0..scale_n {
242                gpu.copy_d2d(
243                    DevicePtr(src.row_scale.0 + (r * s_src_row_bytes + s_col_offset) as u64),
244                    DevicePtr(s_dst.0 + (r * s_local_row_bytes) as u64),
245                    s_local_row_bytes,
246                )?;
247            }
248            Ok(Fp8Weight {
249                weight: w_dst,
250                row_scale: s_dst,
251                n: src.n,
252                k: local_k as u32,
253                scale_format: src.scale_format,
254            })
255        }
256    }
257}
spark_model/tp_shard/quant_shard.rs

spark_model/tp_shard/
quant_shard.rs
