spark_model/layers/ngram_embed/
table.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! One n-gram lookup table on device: BF16 as shipped, FP8-quantized at
4//! load, or NVMe-backed by a bounded row cache.
5
6use anyhow::Result;
7use spark_runtime::gpu::GpuBackend;
8
9use crate::weight_map::{DenseWeight, Fp8DenseWeight};
10
11/// One n-gram lookup table on device: BF16 as shipped, or FP8-quantized at
12/// load (per-row E4M3 + f32 scale — halves the ~63 GB table footprint;
13/// embeddings tolerate this well and the gather dequantizes on read).
14pub enum NgramTable {
15    Bf16(DenseWeight),
16    Fp8(Fp8DenseWeight),
17    /// NVMe-backed: only a bounded set of ROWS is resident, in a pinned
18    /// GPU-addressable arena. The host resolves `row_id -> slot` (the ids are
19    /// a pure function of token ids, so this is host-side anyway) and the
20    /// SAME gather kernels then read the arena by slot index — no kernel
21    /// change, no `cuMemcpyHtoD` on the fault path.
22    ///
23    /// This is what makes a 51 B-parameter embedding table serveable on a
24    /// 121 GB box: the tables are the model's largest tensors and its least
25    /// bandwidth-hungry (12 rows ~ 3 KB per token), so demoting them buys
26    /// back tens of GB for KV.
27    #[cfg(feature = "cuda")]
28    Cached(Box<spark_storage::NgramRowCache>),
29}
30
31impl NgramTable {
32    /// Quantize a BF16 table to FP8 on the GPU (per-row E4M3 absmax +
33    /// f32 scale via `quantize_bf16_to_fp8`) — the quantize-on-load
34    /// lever. The caller frees the BF16 source afterwards; tables are
35    /// loaded one at a time so peak overhead is a single table.
36    pub fn quantize_bf16(
37        w: &DenseWeight,
38        rows: usize,
39        dim: usize,
40        gpu: &dyn GpuBackend,
41        stream: u64,
42    ) -> Result<Self> {
43        let k = gpu.kernel("gemv_fp8w", "quantize_bf16_to_fp8")?;
44        Ok(Self::Fp8(crate::weight_map::quantize_to_fp8(
45            w, rows, dim, gpu, k, stream,
46        )?))
47    }
48}
spark_model/layers/ngram_embed/table.rs

spark_model/layers/ngram_embed/
table.rs
