metrale_model_layers/layers/ngram_embed/
table.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: One n-gram lookup table on device: BF16 as shipped,
4//! FP8-quantized at load, or NVMe-backed by a bounded row cache.
5//!
6//! Owner: model-layers (n-gram embedding).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::GpuBackend;
11
12use crate::weight_map::{DenseWeight, Fp8DenseWeight};
13
14/// 2026-09-25: One n-gram lookup table. `Fp8` holds E4M3 rows with a
15/// per-row f32 scale, which the gather (`batched_embed_fp8`) applies.
16pub enum NgramTable {
17    Bf16(DenseWeight),
18    Fp8(Fp8DenseWeight),
19    /// 2026-09-25: NVMe-backed: a bounded set of rows is resident in a
20    /// GPU-addressable arena. The host maps row ids to arena slots
21    /// (`NgramRowCache::resolve`) and the same gather kernels read the arena
22    /// by slot.
23    #[cfg(feature = "cuda")]
24    Cached(Box<metrale_storage::NgramRowCache>),
25}
26
27impl NgramTable {
28    /// 2026-09-25: Quantize a BF16 table to FP8 on the GPU with
29    /// `quantize_bf16_to_fp8`. `w` is not freed.
30    pub fn quantize_bf16(
31        w: &DenseWeight,
32        rows: usize,
33        dim: usize,
34        gpu: &dyn GpuBackend,
35        stream: u64,
36    ) -> Result<Self> {
37        let k = gpu.kernel("gemv_fp8w", "quantize_bf16_to_fp8")?;
38        Ok(Self::Fp8(crate::weight_map::quantize_to_fp8(
39            w, rows, dim, gpu, k, stream,
40        )?))
41    }
42}
metrale_model_layers/layers/ngram_embed/table.rs

metrale_model_layers/layers/ngram_embed/
table.rs
