metrale_cache/
kv_dequant.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host-side dequantization of paged KV cache blocks to BF16.
4//!
5//! The `--high-speed-swap` offload in `metrale-model-layers`
6//! (`qwen3_attention/decode/high_speed_swap.rs`) uses it to turn quantized K
7//! and V blocks into the BF16 it offloads. The layouts follow the kernels
8//! (`kernels/gb10/common/`):
9//!
10//! | Quant   | data bytes/elem | scale             | LUT                         | kernel                                           |
11//! |---------|-----------------|-------------------|-----------------------------|--------------------------------------------------|
12//! | FP8     | 1 (E4M3)        | per tensor        | `e4m3_lut`                  | `reshape_and_cache_flash_fp8` (reshape_and_cache.cu) |
13//! | NVFP4   | 0.5 (4-bit)     | FP8 per group     | `NVFP4_E2M1_LUT`            | paged_decode_attn_nvfp4.cu                       |
14//! | Turbo4  | 0.5 (4-bit)     | FP8 per group     | `TURBO4_LUT` (16 levels)    | paged_decode_attn_turbo4.cu                      |
15//! | Turbo3  | 0.375 (3-bit)   | FP8 per group     | `TURBO3_LUT` (8 levels)     | paged_decode_attn_turbo3.cu                      |
16//! | Turbo8  | 1 (E4M3)        | BF16 per group    | `e4m3_lut`                  | paged_decode_attn_turbo8.cu                      |
17//!
18//! A group is `NVFP4_GROUP_SIZE` (16) elements. Group scales sit in a section
19//! after the block's data section.
20//!
21//! Owner: cache.
22//! Invariants: none beyond the types.
23
24use half::bf16;
25
26#[path = "kv_dequant/luts.rs"]
27mod luts;
28pub use luts::{NVFP4_E2M1_LUT, NVFP4_GROUP_SIZE, TURBO3_LUT, TURBO4_LUT, e4m3_lut};
29
30/// 2026-09-25: FP8 E4M3 bytes to BF16, times a per-tensor scale.
31pub fn dequant_fp8_to_bf16(fp8_bytes: &[u8], scale: f32, out: &mut [bf16]) {
32    debug_assert_eq!(fp8_bytes.len(), out.len());
33    let lut = e4m3_lut();
34    for (i, b) in fp8_bytes.iter().enumerate() {
35        out[i] = bf16::from_f32(lut[*b as usize] * scale);
36    }
37}
38
39/// 2026-09-25: A 4-bit block (NVFP4 or Turbo4, chosen by `lut`) to BF16.
40///
41/// Layout per block:
42///   data:   `bs * nkv * (hd / 2)` bytes, low nibble first.
43///   scales: `bs * nkv * (hd / NVFP4_GROUP_SIZE)` bytes, one FP8 scale per group.
44pub fn dequant_4bit_block_to_bf16(
45    raw: &[u8],
46    bs: usize,
47    nkv: usize,
48    hd: usize,
49    lut: &[f32; 16],
50    out: &mut [bf16],
51) {
52    debug_assert!(hd.is_multiple_of(NVFP4_GROUP_SIZE));
53    debug_assert!(hd.is_multiple_of(2));
54    debug_assert_eq!(out.len(), bs * nkv * hd);
55    let head_data_bytes = hd / 2;
56    let head_scale_bytes = hd / NVFP4_GROUP_SIZE;
57    let token_data_stride = nkv * head_data_bytes;
58    let token_scale_stride = nkv * head_scale_bytes;
59    let data_section_bytes = bs * token_data_stride;
60    debug_assert!(raw.len() >= data_section_bytes + bs * token_scale_stride);
61    let (data, scales) = raw.split_at(data_section_bytes);
62    let e4m3 = e4m3_lut();
63    for tok in 0..bs {
64        for kv_h in 0..nkv {
65            let d_off = tok * token_data_stride + kv_h * head_data_bytes;
66            let s_off = tok * token_scale_stride + kv_h * head_scale_bytes;
67            for byte_idx in 0..head_data_bytes {
68                let byte = data[d_off + byte_idx];
69                let n0 = (byte & 0x0F) as usize;
70                let n1 = ((byte >> 4) & 0x0F) as usize;
71                let elem_pair_idx = byte_idx * 2;
72                let group_idx = elem_pair_idx / NVFP4_GROUP_SIZE;
73                let scale = e4m3[scales[s_off + group_idx] as usize];
74                let v0 = lut[n0] * scale;
75                let v1 = lut[n1] * scale;
76                let out_base = (tok * nkv + kv_h) * hd + elem_pair_idx;
77                out[out_base] = bf16::from_f32(v0);
78                out[out_base + 1] = bf16::from_f32(v1);
79            }
80        }
81    }
82}
83
84/// 2026-09-25: A Turbo3 block to BF16.
85///
86/// Layout per block:
87///   data:   `bs * nkv * (hd * 3 / 8)` bytes, 8 values in 3 bytes.
88///   scales: `bs * nkv * (hd / NVFP4_GROUP_SIZE)` bytes.
89/// Unpacking of v0..v7 from b0, b1, b2, as in `nvfp4_dequant` of
90/// `kernels/gb10/common/paged_decode_attn_turbo3.cu`:
91///   v0 = b0 & 0x7
92///   v1 = (b0 >> 3) & 0x7
93///   v2 = ((b0 >> 6) | (b1 << 2)) & 0x7
94///   v3 = (b1 >> 1) & 0x7
95///   v4 = (b1 >> 4) & 0x7
96///   v5 = ((b1 >> 7) | (b2 << 1)) & 0x7
97///   v6 = (b2 >> 2) & 0x7
98///   v7 = (b2 >> 5) & 0x7
99pub fn dequant_turbo3_block_to_bf16(
100    raw: &[u8],
101    bs: usize,
102    nkv: usize,
103    hd: usize,
104    out: &mut [bf16],
105) {
106    debug_assert!(hd.is_multiple_of(8));
107    debug_assert!(hd.is_multiple_of(NVFP4_GROUP_SIZE));
108    debug_assert_eq!(out.len(), bs * nkv * hd);
109    let head_data_bytes = hd * 3 / 8;
110    let head_scale_bytes = hd / NVFP4_GROUP_SIZE;
111    let token_data_stride = nkv * head_data_bytes;
112    let token_scale_stride = nkv * head_scale_bytes;
113    let data_section_bytes = bs * token_data_stride;
114    debug_assert!(raw.len() >= data_section_bytes + bs * token_scale_stride);
115    let (data, scales) = raw.split_at(data_section_bytes);
116    let e4m3 = e4m3_lut();
117    for tok in 0..bs {
118        for kv_h in 0..nkv {
119            let d_off = tok * token_data_stride + kv_h * head_data_bytes;
120            let s_off = tok * token_scale_stride + kv_h * head_scale_bytes;
121            for triplet_idx in 0..hd / 8 {
122                let b0 = data[d_off + triplet_idx * 3] as u32;
123                let b1 = data[d_off + triplet_idx * 3 + 1] as u32;
124                let b2 = data[d_off + triplet_idx * 3 + 2] as u32;
125                let nibbles = [
126                    (b0) & 0x7,
127                    (b0 >> 3) & 0x7,
128                    ((b0 >> 6) | (b1 << 2)) & 0x7,
129                    (b1 >> 1) & 0x7,
130                    (b1 >> 4) & 0x7,
131                    ((b1 >> 7) | (b2 << 1)) & 0x7,
132                    (b2 >> 2) & 0x7,
133                    (b2 >> 5) & 0x7,
134                ];
135                let elem_base_in_head = triplet_idx * 8;
136                for k in 0..8 {
137                    let elem = elem_base_in_head + k;
138                    let group_idx = elem / NVFP4_GROUP_SIZE;
139                    let scale = e4m3[scales[s_off + group_idx] as usize];
140                    let v = TURBO3_LUT[nibbles[k] as usize] * scale;
141                    let out_idx = (tok * nkv + kv_h) * hd + elem;
142                    out[out_idx] = bf16::from_f32(v);
143                }
144            }
145        }
146    }
147}
148
149/// 2026-09-25: A Turbo8 block (FP8 E4M3 data, BF16 group scales) to BF16.
150///
151/// Layout per block:
152///   data:   `bs * nkv * hd` bytes, one FP8 byte per element.
153///   scales: `bs * nkv * (hd / NVFP4_GROUP_SIZE) * 2` bytes, little-endian BF16.
154/// The scale read matches `nvfp4_dequant` of
155/// `kernels/gb10/common/paged_decode_attn_turbo8.cu`.
156pub fn dequant_turbo8_block_to_bf16(
157    raw: &[u8],
158    bs: usize,
159    nkv: usize,
160    hd: usize,
161    out: &mut [bf16],
162) {
163    debug_assert!(hd.is_multiple_of(NVFP4_GROUP_SIZE));
164    debug_assert_eq!(out.len(), bs * nkv * hd);
165    let head_data_bytes = hd;
166    let head_scale_bytes = (hd / NVFP4_GROUP_SIZE) * 2;
167    let token_data_stride = nkv * head_data_bytes;
168    let token_scale_stride = nkv * head_scale_bytes;
169    let data_section_bytes = bs * token_data_stride;
170    debug_assert!(raw.len() >= data_section_bytes + bs * token_scale_stride);
171    let (data, scales) = raw.split_at(data_section_bytes);
172    let e4m3 = e4m3_lut();
173    for tok in 0..bs {
174        for kv_h in 0..nkv {
175            let d_off = tok * token_data_stride + kv_h * head_data_bytes;
176            let s_off = tok * token_scale_stride + kv_h * head_scale_bytes;
177            for i in 0..hd {
178                let byte = data[d_off + i] as usize;
179                let group_idx = i / NVFP4_GROUP_SIZE;
180                let s_byte_off = s_off + group_idx * 2;
181                let scale_bf16 = bf16::from_le_bytes([scales[s_byte_off], scales[s_byte_off + 1]]);
182                let scale = scale_bf16.to_f32();
183                let v = e4m3[byte] * scale;
184                let out_idx = (tok * nkv + kv_h) * hd + i;
185                out[out_idx] = bf16::from_f32(v);
186            }
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    /// 2026-09-25: FP8 E4M3 +1.0: sign 0, exponent 7 (bias 7), mantissa 0.
196    const FP8_ONE: u8 = 0x38;
197
198    #[test]
199    fn e4m3_lut_basics() {
200        let lut = e4m3_lut();
201        assert_eq!(lut[0x00], 0.0);
202        assert_eq!(lut[0x80], -0.0);
203        assert!((lut[0x38] - 1.0).abs() < 1e-6);
204        assert!((lut[0xB8] + 1.0).abs() < 1e-6);
205        assert!((lut[0x3F] - 1.875).abs() < 1e-6);
206        assert!((lut[0x40] - 2.0).abs() < 1e-6);
207        assert!((lut[0x78] - 256.0).abs() < 1e-6);
208        assert!((lut[0x7E] - 448.0).abs() < 1e-6);
209        assert!(lut[0x7F].is_nan());
210        assert!(lut[0xFF].is_nan());
211        assert!((lut[0x01] - (1.0 / 512.0)).abs() < 1e-9);
212    }
213
214    #[test]
215    fn fp8_dequant_with_scale() {
216        let bytes = [0x38u8, 0x3F, 0x40, 0xB8];
217        let mut out = vec![bf16::ZERO; 4];
218        dequant_fp8_to_bf16(&bytes, 0.5, &mut out);
219        let f: Vec<f32> = out.iter().map(|x| x.to_f32()).collect();
220        assert!((f[0] - 0.5).abs() < 1e-2);
221        assert!((f[1] - 0.9375).abs() < 1e-2);
222        assert!((f[2] - 1.0).abs() < 1e-2);
223        assert!((f[3] + 0.5).abs() < 1e-2);
224    }
225
226    #[test]
227    #[allow(clippy::needless_range_loop)]
228    fn nvfp4_dequant_layout() {
229        let bs = 1;
230        let nkv = 1;
231        let hd = 16;
232        let mut raw = vec![0u8; 8 + 1];
233        for i in 0..8 {
234            let lo = (2 * i) & 0xF;
235            let hi = (2 * i + 1) & 0xF;
236            raw[i] = (lo as u8) | ((hi as u8) << 4);
237        }
238        raw[8] = FP8_ONE;
239        let mut out = vec![bf16::ZERO; bs * nkv * hd];
240        dequant_4bit_block_to_bf16(&raw, bs, nkv, hd, &NVFP4_E2M1_LUT, &mut out);
241        for i in 0..hd {
242            let expected = NVFP4_E2M1_LUT[i];
243            assert!(
244                (out[i].to_f32() - expected).abs() < 1e-2,
245                "elem {i}: expected {expected}, got {}",
246                out[i].to_f32(),
247            );
248        }
249    }
250
251    #[test]
252    #[allow(clippy::needless_range_loop)]
253    fn turbo4_dequant_layout() {
254        // 2026-09-25: The NVFP4 test's nibbles with `TURBO4_LUT`: the LUT
255        // argument is used.
256        let bs = 1;
257        let nkv = 1;
258        let hd = 16;
259        let mut raw = vec![0u8; 8 + 1];
260        for i in 0..8 {
261            raw[i] = ((2 * i) as u8 & 0xF) | (((2 * i + 1) as u8 & 0xF) << 4);
262        }
263        raw[8] = FP8_ONE;
264        let mut out = vec![bf16::ZERO; hd];
265        dequant_4bit_block_to_bf16(&raw, bs, nkv, hd, &TURBO4_LUT, &mut out);
266        for i in 0..hd {
267            assert!(
268                (out[i].to_f32() - TURBO4_LUT[i]).abs() < 1e-2,
269                "elem {i}: expected {}, got {}",
270                TURBO4_LUT[i],
271                out[i].to_f32(),
272            );
273        }
274    }
275
276    #[test]
277    fn turbo3_unpack_round_trip() {
278        let bs = 1;
279        let nkv = 1;
280        let hd = 16;
281        let head_data_bytes = hd * 3 / 8;
282        let mut raw = vec![0u8; head_data_bytes + 1];
283        let pack8 = |vals: [u8; 8]| -> [u8; 3] {
284            let b0 = vals[0] | (vals[1] << 3) | (vals[2] << 6);
285            let b1 = (vals[2] >> 2) | (vals[3] << 1) | (vals[4] << 4) | (vals[5] << 7);
286            let b2 = (vals[5] >> 1) | (vals[6] << 2) | (vals[7] << 5);
287            [b0, b1, b2]
288        };
289        let t0 = pack8([0, 1, 2, 3, 4, 5, 6, 7]);
290        let t1 = pack8([7, 6, 5, 4, 3, 2, 1, 0]);
291        raw[..3].copy_from_slice(&t0);
292        raw[3..6].copy_from_slice(&t1);
293        raw[head_data_bytes] = FP8_ONE;
294        let mut out = vec![bf16::ZERO; bs * nkv * hd];
295        dequant_turbo3_block_to_bf16(&raw, bs, nkv, hd, &mut out);
296        let expect: Vec<f32> = (0..8u32)
297            .map(|i| TURBO3_LUT[i as usize])
298            .chain((0..8u32).rev().map(|i| TURBO3_LUT[i as usize]))
299            .collect();
300        for (i, e) in expect.iter().enumerate() {
301            assert!(
302                (out[i].to_f32() - e).abs() < 1e-2,
303                "elem {i}: expected {e}, got {}",
304                out[i].to_f32(),
305            );
306        }
307    }
308
309    #[test]
310    #[allow(clippy::needless_range_loop)]
311    fn turbo8_dequant_layout() {
312        // 2026-09-25: One token, one KV head, hd=16: one group of 16 FP8 bytes, then its
313        // BF16 scale (2 bytes, not a 1-byte FP8 scale). 18 bytes in all.
314        let bs = 1;
315        let nkv = 1;
316        let hd = 16;
317        let mut raw = vec![FP8_ONE; hd + 2];
318        let scale_bytes = bf16::from_f32(1.0).to_le_bytes();
319        raw[hd] = scale_bytes[0];
320        raw[hd + 1] = scale_bytes[1];
321        let mut out = vec![bf16::ZERO; bs * nkv * hd];
322        dequant_turbo8_block_to_bf16(&raw, bs, nkv, hd, &mut out);
323        for i in 0..hd {
324            assert!(
325                (out[i].to_f32() - 1.0).abs() < 1e-2,
326                "elem {i}: expected 1.0, got {}",
327                out[i].to_f32(),
328            );
329        }
330    }
331
332    #[test]
333    fn multi_head_multi_token_consistency() {
334        // 2026-09-25: 2 tokens x 2 KV heads x hd 16: four (token, head) groups
335        // of 8 data bytes and 1 scale byte, each with its own nibble pair.
336        let bs = 2;
337        let nkv = 2;
338        let hd = 16;
339        let head_data_bytes = hd / 2;
340        let head_scale_bytes = hd / NVFP4_GROUP_SIZE;
341        let token_data_stride = nkv * head_data_bytes;
342        let token_scale_stride = nkv * head_scale_bytes;
343        let data_section_bytes = bs * token_data_stride;
344        let scale_section_bytes = bs * token_scale_stride;
345        let mut raw = vec![0u8; data_section_bytes + scale_section_bytes];
346        for tok in 0..bs {
347            for kv_h in 0..nkv {
348                let d_off = tok * token_data_stride + kv_h * head_data_bytes;
349                let nibble_lo = ((tok * 2 + kv_h) as u8) & 0xF;
350                let nibble_hi = (((tok * 2 + kv_h) + 8) as u8) & 0xF;
351                for byte_idx in 0..head_data_bytes {
352                    raw[d_off + byte_idx] = nibble_lo | (nibble_hi << 4);
353                }
354                let s_off = data_section_bytes + tok * token_scale_stride + kv_h * head_scale_bytes;
355                raw[s_off] = FP8_ONE;
356            }
357        }
358        let mut out = vec![bf16::ZERO; bs * nkv * hd];
359        dequant_4bit_block_to_bf16(&raw, bs, nkv, hd, &NVFP4_E2M1_LUT, &mut out);
360        for tok in 0..bs {
361            for kv_h in 0..nkv {
362                let nibble_lo = (tok * 2 + kv_h) & 0xF;
363                let nibble_hi = ((tok * 2 + kv_h) + 8) & 0xF;
364                let exp_lo = NVFP4_E2M1_LUT[nibble_lo];
365                let exp_hi = NVFP4_E2M1_LUT[nibble_hi];
366                let base = (tok * nkv + kv_h) * hd;
367                for elem in 0..hd {
368                    let expected = if elem % 2 == 0 { exp_lo } else { exp_hi };
369                    assert!(
370                        (out[base + elem].to_f32() - expected).abs() < 1e-2,
371                        "tok={tok} kv_h={kv_h} elem={elem}: expected {expected}, got {}",
372                        out[base + elem].to_f32(),
373                    );
374                }
375            }
376        }
377    }
378}
metrale_cache/kv_dequant.rs

metrale_cache/
kv_dequant.rs
