diff --git a/crates/larql-compute/csrc/q4_dot.c b/crates/larql-compute/csrc/q4_dot.c index e27c3eab2..8db120a75 100644 --- a/crates/larql-compute/csrc/q4_dot.c +++ b/crates/larql-compute/csrc/q4_dot.c @@ -70,13 +70,12 @@ float q4_q8_dot_neon_c( // Load 16 bytes of packed Q4 nibbles uint8x16_t raw = vld1q_u8(quants); - // Split into low/high nibbles, subtract 8 for signed range - int8x16_t lo = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(raw, mask_lo)), offset); - int8x16_t hi = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(raw, 4)), offset); - - // Interleave: [lo0,hi0,lo1,hi1,...] to match sequential Q8 layout - int8x16_t q4_0 = vzip1q_s8(lo, hi); // first 16 interleaved values - int8x16_t q4_1 = vzip2q_s8(lo, hi); // next 16 interleaved values + // Split into low/high nibbles, subtract 8 for signed range. + // ggml planar Q4_0 layout (quantize_row_q4_0_ref): low nibbles are + // elements 0..16, high nibbles are elements 16..32 — they pair with + // the sequential Q8 input directly, no interleave needed. + int8x16_t q4_0 = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(raw, mask_lo)), offset); + int8x16_t q4_1 = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(raw, 4)), offset); // Load Q8 values int8x16_t q8_0 = vld1q_s8(q8_ptr); @@ -157,8 +156,9 @@ void q4_0_vecmat_c( uint8_t byte = quants[j]; int lo_v = (byte & 0x0F) - 8; int hi_v = ((byte >> 4) & 0x0F) - 8; - o[j * 2] += (float)lo_v * scale; - o[j * 2 + 1] += (float)hi_v * scale; + // ggml planar layout: lo → element j, hi → element j+16. + o[j] += (float)lo_v * scale; + o[j + 16] += (float)hi_v * scale; } } } @@ -191,8 +191,9 @@ void q4_0_matvec_c( uint8_t byte = quants[j]; int lo_v = (byte & 0x0F) - 8; int hi_v = ((byte >> 4) & 0x0F) - 8; - acc += (float)lo_v * (float)q8_ptr[j * 2] * combined_scale; - acc += (float)hi_v * (float)q8_ptr[j * 2 + 1] * combined_scale; + // ggml planar layout: lo → element j, hi → element j+16. + acc += (float)lo_v * (float)q8_ptr[j] * combined_scale; + acc += (float)hi_v * (float)q8_ptr[j + 16] * combined_scale; } } scores[row] = acc; @@ -225,8 +226,9 @@ void q4_0_vecmat_c( uint8_t byte = quants[j]; int lo_v = (byte & 0x0F) - 8; int hi_v = ((byte >> 4) & 0x0F) - 8; - o[j * 2] += (float)lo_v * scale; - o[j * 2 + 1] += (float)hi_v * scale; + // ggml planar layout: lo → element j, hi → element j+16. + o[j] += (float)lo_v * scale; + o[j + 16] += (float)hi_v * scale; } } } diff --git a/crates/larql-compute/src/cpu/ops/q4_common.rs b/crates/larql-compute/src/cpu/ops/q4_common.rs index 4912819c2..0c91860c3 100644 --- a/crates/larql-compute/src/cpu/ops/q4_common.rs +++ b/crates/larql-compute/src/cpu/ops/q4_common.rs @@ -81,9 +81,11 @@ pub fn quantize_q4_0(data: &[f32]) -> Vec { } }; out.extend_from_slice(&f16.to_le_bytes()); + // ggml planar nibble layout (`quantize_row_q4_0_ref`): byte j packs + // element j (low nibble) and element j+16 (high nibble). for j in 0..16 { - let lo = ((block[j * 2] * inv).round() as i32 + 8).clamp(0, 15) as u8; - let hi = ((block[j * 2 + 1] * inv).round() as i32 + 8).clamp(0, 15) as u8; + let lo = ((block[j] * inv).round() as i32 + 8).clamp(0, 15) as u8; + let hi = ((block[j + 16] * inv).round() as i32 + 8).clamp(0, 15) as u8; out.push(lo | (hi << 4)); } } @@ -303,21 +305,29 @@ pub fn quantize_q6_k(data: &[f32]) -> Vec { } } - // Pack lower 4 bits: 128 bytes (2 nibbles per byte) + // Pack per ggml's planar Q6_K layout (`quantize_row_q6_K_ref`): + // within each 128-element half, ql[l] holds element l in its low + // nibble and element l+64 in its high nibble; ql[l+32] holds + // elements l+32 / l+96. qh[l] packs the two high bits of elements + // l, l+32, l+64, l+96 at shifts 0/2/4/6. let mut ql = [0u8; 128]; - for i in 0..128 { - ql[i] = (q6_vals[i * 2] & 0x0F) | ((q6_vals[i * 2 + 1] & 0x0F) << 4); - } - out.extend_from_slice(&ql); - - // Pack upper 2 bits: 64 bytes (4 × 2 bits per byte) let mut qh = [0u8; 64]; - for (i, &q6_val) in q6_vals.iter().enumerate() { - let hi2 = (q6_val >> 4) & 0x03; - let byte_idx = i / 4; - let bit_offset = (i % 4) * 2; - qh[byte_idx] |= hi2 << bit_offset; + for half in 0..2 { + let e = half * 128; // element base for this half + for l in 0..32 { + let q1 = q6_vals[e + l]; + let q2 = q6_vals[e + l + 32]; + let q3 = q6_vals[e + l + 64]; + let q4 = q6_vals[e + l + 96]; + ql[half * 64 + l] = (q1 & 0x0F) | ((q3 & 0x0F) << 4); + ql[half * 64 + l + 32] = (q2 & 0x0F) | ((q4 & 0x0F) << 4); + qh[half * 32 + l] = ((q1 >> 4) & 3) + | (((q2 >> 4) & 3) << 2) + | (((q3 >> 4) & 3) << 4) + | (((q4 >> 4) & 3) << 6); + } } + out.extend_from_slice(&ql); out.extend_from_slice(&qh); // 16 × int8 scales @@ -706,22 +716,13 @@ fn decode_q4k_superblock_into(w: &[u8], row_base: usize, sb: usize, wf: &mut [f3 fn decode_q6k_superblock_into(w: &[u8], row_base: usize, sb: usize, wf: &mut [f32; 256]) { const BLOCK_BYTES: usize = 210; let block = &w[row_base + sb * BLOCK_BYTES..row_base + (sb + 1) * BLOCK_BYTES]; - let ql = &block[0..128]; - let qh = &block[128..192]; let scales = &block[192..208]; let d = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); for (j, &sc_byte) in scales.iter().enumerate() { let sc = d * (sc_byte as i8) as f32; - for i in 0..16 { - let idx = j * 16 + i; - let lo4 = if idx % 2 == 0 { - ql[idx / 2] & 0x0F - } else { - (ql[idx / 2] >> 4) & 0x0F - }; - let hi2 = (qh[idx / 4] >> ((idx % 4) * 2)) & 0x03; - let val = ((lo4 as i32) | ((hi2 as i32) << 4)) - 32; - wf[idx] = sc * val as f32; + let vals = larql_models::quant::ggml::q6_k::q6k_subblock_vals(block, j); + for (i, &v) in vals.iter().enumerate() { + wf[j * 16 + i] = sc * v as f32; } } } @@ -1690,13 +1691,12 @@ mod tests { let scale_bits = u16::from_le_bytes([q4[0], q4[1]]); let scale = f16_to_f32(scale_bits); - let mut decoded = Vec::with_capacity(32); + // ggml planar layout: low nibbles are elements 0..16, high 16..32. + let mut decoded = vec![0.0f32; 32]; for j in 0..16 { let byte = q4[2 + j]; - let lo = (byte & 0x0F) as i32 - 8; - let hi = (byte >> 4) as i32 - 8; - decoded.push(lo as f32 * scale); - decoded.push(hi as f32 * scale); + decoded[j] = ((byte & 0x0F) as i32 - 8) as f32 * scale; + decoded[j + 16] = ((byte >> 4) as i32 - 8) as f32 * scale; } // Check approximate reconstruction (Q4 is lossy, but should be close) diff --git a/crates/larql-compute/src/cpu/ops/q4k_q8k_dot.rs b/crates/larql-compute/src/cpu/ops/q4k_q8k_dot.rs index d5aab0873..2bbed8b61 100644 --- a/crates/larql-compute/src/cpu/ops/q4k_q8k_dot.rs +++ b/crates/larql-compute/src/cpu/ops/q4k_q8k_dot.rs @@ -1746,7 +1746,10 @@ pub fn q4k_q8k_gate_up_asm( // [192..208] 16 bytes: scales — one int8 per 16 elements // [208..210] 2 bytes: d — f16 super-block scale // -// Element i: raw6 = (ql[i/2] >> 4*(i&1)) & 0xF | (((qh[i/4] >> 2*(i%4)) & 3) << 4) +// Element placement follows ggml's planar layout (see +// `larql_models::quant::ggml::q6_k::q6k_subblock_vals`): within each +// 128-element half, ql low nibbles hold elements 0..63 and high nibbles +// 64..127; qh[l] packs the hi2 bits of elements l/l+32/l+64/l+96. // w[i] = d * scales[i/16] * (raw6 - 32) // // Dot product with Q8_K activation `q8k`: @@ -1782,8 +1785,6 @@ pub fn q6k_q8k_matvec_scalar( let mut acc = 0.0f32; for sb in 0..n_blocks { let block = &w[row_base + sb * Q6K_BLOCK_BYTES..]; - let ql = &block[0..128]; - let qh = &block[128..192]; let sc = &block[192..208]; // 16 × int8 let d_w = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); let d_y = q8k_x.d[sb]; @@ -1792,20 +1793,16 @@ pub fn q6k_q8k_matvec_scalar( let mut sum1: i32 = 0; for (g, scale_byte) in sc.iter().enumerate().take(16usize) { - // 16-element group g, using scale sc[g]. + // 16-element group g, using scale sc[g]. Weights decode + // through the shared ggml planar-layout helper. let scale = *scale_byte as i8 as i32; + let vals = larql_models::quant::ggml::q6_k::q6k_subblock_vals( + &block[..Q6K_BLOCK_BYTES], + g, + ); let mut dot_g: i32 = 0; - for k in 0..16usize { - let i = g * 16 + k; - let lo4 = if i & 1 == 0 { - (ql[i / 2] & 0x0F) as i32 - } else { - ((ql[i / 2] >> 4) & 0x0F) as i32 - }; - let hi2 = ((qh[i / 4] >> (2 * (i % 4))) & 0x03) as i32; - let raw6 = lo4 | (hi2 << 4); - let w_i = raw6 - 32; - dot_g += w_i * q8_qs[i] as i32; + for (k, &v) in vals.iter().enumerate() { + dot_g += (v as i32) * q8_qs[g * 16 + k] as i32; } sum1 += scale * dot_g; } @@ -1825,6 +1822,11 @@ pub fn q6k_q8k_matvec_scalar( /// 3. scale * dot_g accumulated into sum1. /// /// Final: acc += d_w * d_y * sum1. +/// +/// TODO(q6k-planar): still decodes the pre-fix interleaved layout — +/// unreachable from the dispatcher until reworked for ggml's planar +/// layout and re-verified on ARM. +#[allow(dead_code)] #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] pub fn q6k_q8k_matvec_neon( out: &mut [f32], @@ -2048,6 +2050,11 @@ unsafe fn q6k_sb_sum1_asm(ql: *const u8, qh: *const u8, act: *const i8, scales: /// epilogue (`acc += d_w·d_y·sum1`, no mins term) is the same Rust code, so /// it is bit-exact with the scalar reference /// (`q6k_matvec_asm_matches_scalar_bit_exact`). +/// +/// TODO(q6k-planar): still decodes the pre-fix interleaved layout — +/// unreachable from the dispatcher until reworked for ggml's planar +/// layout and re-verified on ARM. +#[allow(dead_code)] #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] pub fn q6k_q8k_matvec_asm( out: &mut [f32], @@ -2105,18 +2112,10 @@ pub fn q6k_q8k_matvec_into( rows: usize, cols: usize, ) { - #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] - { - // C12: same opt-in as the Q4_K kernels — `LARQL_Q4K_ASM=1` routes - // through the hand-asm form. Bit-exact; default off. - if use_asm_kernel() { - q6k_q8k_matvec_asm(out, q8k_x, w, rows, cols); - } else { - q6k_q8k_matvec_neon(out, q8k_x, w, rows, cols); - } - return; - } - #[allow(unreachable_code)] + // TODO(q6k-planar): the NEON and hand-asm forms still decode the + // pre-fix interleaved nibble layout; they need the same ggml-planar + // rework as the scalar path (and verification on ARM hardware) before + // they can be re-enabled. Until then every arch takes the scalar path. q6k_q8k_matvec_scalar(out, q8k_x, w, rows, cols); } diff --git a/crates/larql-compute/src/cpu/ops/q6k_matvec.rs b/crates/larql-compute/src/cpu/ops/q6k_matvec.rs index dda83de4f..31509021c 100644 --- a/crates/larql-compute/src/cpu/ops/q6k_matvec.rs +++ b/crates/larql-compute/src/cpu/ops/q6k_matvec.rs @@ -1,8 +1,11 @@ //! CPU reference implementation for Q6_K matrix-vector multiply. //! -//! Mirrors the Metal shader `q6k_matvec` exactly for cross-backend testing. -//! Not optimised — scalar code intended as a correctness reference. +//! Decodes ggml's planar Q6_K super-block layout through the shared +//! `larql_models::quant::ggml::q6_k::q6k_subblock_vals` helper — the +//! single source of truth for Q6_K bit placement. Not optimised — +//! scalar code intended as a correctness reference. +use larql_models::quant::ggml::q6_k::q6k_subblock_vals; use larql_models::quant::ggml::Q6_K_BLOCK_BYTES as Q6K_BLOCK_SIZE; /// Decode f16 bits to f32. @@ -38,7 +41,7 @@ fn f16_to_f32(bits: u16) -> f32 { /// CPU Q6_K matvec: out[N] = Q6_K[N, K] @ x[K]. /// -/// Mirrors the Metal `q6k_matvec` shader: per-row dot product over super-blocks. +/// Per-row dot product over super-blocks, decoded in ggml's planar layout. pub fn dispatch(q6k_data: &[u8], x: &[f32], num_rows: usize, hidden: usize) -> Vec { let superblocks = hidden / 256; let bytes_per_row = superblocks * Q6K_BLOCK_SIZE; @@ -64,26 +67,22 @@ pub fn dispatch(q6k_data: &[u8], x: &[f32], num_rows: usize, hidden: usize) -> V let mut acc = 0.0f32; for sb in 0..superblocks { - let block = &q6k_ref[row_start + sb * Q6K_BLOCK_SIZE..]; - - let ql = &block[0..128]; - let qh = &block[128..192]; + let block = + &q6k_ref[row_start + sb * Q6K_BLOCK_SIZE..][..Q6K_BLOCK_SIZE]; let scales = &block[192..208]; let d_bits = u16::from_le_bytes([block[208], block[209]]); let d = f16_to_f32(d_bits); - let x_base = sb * 256; for (j, &scale) in scales.iter().enumerate() { let sc = d * (scale as i8) as f32; - // Sub-block of 16 elements: ql[j*8 .. j*8+8] gives 16 - // 4-bit lo values; qh[j*4 .. j*4+4] gives 16 2-bit hi - // values (4 packed into each byte). - let ql_sub = &ql[j * 8..j * 8 + 8]; - let qh_sub = &qh[j * 4..j * 4 + 4]; + let vals = q6k_subblock_vals(block, j); let x_sub = &x_ref[x_base + j * 16..x_base + j * 16 + 16]; - - acc += sc * q6_subblock_dot_16(ql_sub, qh_sub, x_sub); + let mut sub = 0.0f32; + for (v, xi) in vals.iter().zip(x_sub) { + sub += *v as f32 * xi; + } + acc += sc * sub; } } *out_val = acc; @@ -92,184 +91,6 @@ pub fn dispatch(q6k_data: &[u8], x: &[f32], num_rows: usize, hidden: usize) -> V out } -/// Decode one 16-element Q6_K sub-block and dot it with 16 f32 inputs. -/// Returns `sum_{i=0..16} ((lo4_i + hi2_i * 16) - 32) * x[i]`. -/// Dispatches to NEON on aarch64; scalar elsewhere. -#[inline] -fn q6_subblock_dot_16(ql_sub: &[u8], qh_sub: &[u8], x_sub: &[f32]) -> f32 { - debug_assert_eq!(ql_sub.len(), 8); - debug_assert_eq!(qh_sub.len(), 4); - debug_assert_eq!(x_sub.len(), 16); - #[cfg(target_arch = "aarch64")] - { - unsafe { q6_subblock_dot_16_neon(ql_sub, qh_sub, x_sub) } - } - #[cfg(not(target_arch = "aarch64"))] - { - let mut acc = 0.0f32; - #[allow(clippy::needless_range_loop)] - for i in 0..8usize { - let qi = i * 2; - let lo_byte = ql_sub[i]; - let lo4_0 = (lo_byte & 0x0F) as f32; - let lo4_1 = ((lo_byte >> 4) & 0x0F) as f32; - let hi_byte_idx_0 = qi / 4; - let hi_byte_idx_1 = (qi + 1) / 4; - let bit_off_0 = (qi % 4) * 2; - let bit_off_1 = ((qi + 1) % 4) * 2; - let hi2_0 = ((qh_sub[hi_byte_idx_0] >> bit_off_0) & 0x03) as f32; - let hi2_1 = ((qh_sub[hi_byte_idx_1] >> bit_off_1) & 0x03) as f32; - let v0 = (lo4_0 + hi2_0 * 16.0) - 32.0; - let v1 = (lo4_1 + hi2_1 * 16.0) - 32.0; - acc += v0 * x_sub[qi] + v1 * x_sub[qi + 1]; - } - acc - } -} - -#[cfg(target_arch = "aarch64")] -#[target_feature(enable = "neon")] -unsafe fn q6_subblock_dot_16_neon(ql_sub: &[u8], qh_sub: &[u8], x_sub: &[f32]) -> f32 { - use core::arch::aarch64::*; - - // ── Low 4 bits: 16 nibbles in output-order ─────────────────────── - // ql_sub is 8 bytes; byte i packs (output 2i, output 2i+1) as - // (lo nibble, hi nibble). Result: u8x16 where lane k = lo4[k]. - let lo_bytes_u64 = u64::from_le_bytes([ - ql_sub[0], ql_sub[1], ql_sub[2], ql_sub[3], ql_sub[4], ql_sub[5], ql_sub[6], ql_sub[7], - ]); - let lo_bytes: uint8x8_t = vcreate_u8(lo_bytes_u64); - let mask4 = vdup_n_u8(0x0F); - let even_lo = vand_u8(lo_bytes, mask4); // lo nibble per byte → outputs 0,2,4,... - let odd_lo = vshr_n_u8::<4>(lo_bytes); // hi nibble per byte → outputs 1,3,5,... - // Interleave to [even[0], odd[0], even[1], odd[1], ...] - let lo16: uint8x16_t = vcombine_u8(vzip1_u8(even_lo, odd_lo), vzip2_u8(even_lo, odd_lo)); - - // ── Hi 2 bits: 16 values, 4 per byte ───────────────────────────── - // qh_sub is 4 bytes; byte i holds outputs (4i, 4i+1, 4i+2, 4i+3) - // at bits (0-1, 2-3, 4-5, 6-7). Broadcast each byte 4× then per- - // lane right-shift by [0, 2, 4, 6] then mask with 0x03. - let qh_bytes_u32 = u32::from_le_bytes([qh_sub[0], qh_sub[1], qh_sub[2], qh_sub[3]]); - // Replicate u32 to fill a u8x16: [b0,b1,b2,b3, b0,b1,b2,b3, ...] - // — we want [b0,b0,b0,b0, b1,b1,b1,b1, b2,b2,b2,b2, b3,b3,b3,b3]. - // tbl with index pattern [0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3]. - let qh_lane: uint8x16_t = vreinterpretq_u8_u32(vdupq_n_u32(qh_bytes_u32)); - #[rustfmt::skip] - let tbl_idx: uint8x16_t = vld1q_u8([ - 0u8, 0, 0, 0, 1, 1, 1, 1, - 2, 2, 2, 2, 3, 3, 3, 3, - ].as_ptr()); - let qh_bcast = vqtbl1q_u8(qh_lane, tbl_idx); - // Per-lane right-shift by [0,2,4,6, 0,2,4,6, ...] using - // vshlq_s8 with negative shifts (treats input as signed s8 but - // we mask immediately after so sign doesn't leak). - #[rustfmt::skip] - let shift_idx: int8x16_t = vld1q_s8([ - 0i8, -2, -4, -6, 0, -2, -4, -6, - 0, -2, -4, -6, 0, -2, -4, -6, - ].as_ptr()); - let hi_shifted = vshlq_u8(qh_bcast, shift_idx); - let mask2 = vdupq_n_u8(0x03); - let hi16 = vandq_u8(hi_shifted, mask2); - - // ── Combine: u8 value = lo4 + hi2 * 16, then -32 in f32 ────────── - // (hi2 << 4) | lo4 fits in u8 (max 63); we widen later. - let combined = vorrq_u8(lo16, vshlq_n_u8::<4>(hi16)); - - // Widen u8x16 → 4× u32x4 → 4× f32x4 and subtract 32. - let lo16u = vmovl_u8(vget_low_u8(combined)); - let hi16u = vmovl_u8(vget_high_u8(combined)); - let v0u = vmovl_u16(vget_low_u16(lo16u)); - let v1u = vmovl_u16(vget_high_u16(lo16u)); - let v2u = vmovl_u16(vget_low_u16(hi16u)); - let v3u = vmovl_u16(vget_high_u16(hi16u)); - let off = vdupq_n_f32(32.0); - let v0 = vsubq_f32(vcvtq_f32_u32(v0u), off); - let v1 = vsubq_f32(vcvtq_f32_u32(v1u), off); - let v2 = vsubq_f32(vcvtq_f32_u32(v2u), off); - let v3 = vsubq_f32(vcvtq_f32_u32(v3u), off); - - // FMA against x_sub[0..16] into four independent accumulators so - // the 4 FMAs pipeline at 1/cycle instead of serialising on a - // single dst register (M3 FMA: 4-cycle latency, 1/cycle throughput). - let x0 = vld1q_f32(x_sub.as_ptr()); - let x1 = vld1q_f32(x_sub.as_ptr().add(4)); - let x2 = vld1q_f32(x_sub.as_ptr().add(8)); - let x3 = vld1q_f32(x_sub.as_ptr().add(12)); - let acc0 = vmulq_f32(v0, x0); - let acc1 = vmulq_f32(v1, x1); - let acc2 = vmulq_f32(v2, x2); - let acc3 = vmulq_f32(v3, x3); - let acc = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3)); - vaddvq_f32(acc) -} - -#[cfg(test)] -mod neon_tests { - use super::*; - - // Reference scalar oracle for the NEON sub-block dot. Indexed - // access mirrors the Q6_K layout walk used by the production - // kernel; switching to enumerate()/iter() obscures the - // sub-block-offset arithmetic that's the point of the test. - #[allow(clippy::needless_range_loop)] - fn scalar_subblock_dot_16(ql_sub: &[u8], qh_sub: &[u8], x_sub: &[f32]) -> f32 { - let mut acc = 0.0f32; - for i in 0..8usize { - let qi = i * 2; - let lo_byte = ql_sub[i]; - let lo4_0 = (lo_byte & 0x0F) as f32; - let lo4_1 = ((lo_byte >> 4) & 0x0F) as f32; - let hi_byte_idx_0 = qi / 4; - let hi_byte_idx_1 = (qi + 1) / 4; - let bit_off_0 = (qi % 4) * 2; - let bit_off_1 = ((qi + 1) % 4) * 2; - let hi2_0 = ((qh_sub[hi_byte_idx_0] >> bit_off_0) & 0x03) as f32; - let hi2_1 = ((qh_sub[hi_byte_idx_1] >> bit_off_1) & 0x03) as f32; - let v0 = (lo4_0 + hi2_0 * 16.0) - 32.0; - let v1 = (lo4_1 + hi2_1 * 16.0) - 32.0; - acc += v0 * x_sub[qi] + v1 * x_sub[qi + 1]; - } - acc - } - - #[test] - fn q6_subblock_matches_scalar_full_6bit_range() { - // Pack lo + hi to cover every 6-bit value 0..63 across 16 - // positions, repeated. - let ql_sub: Vec = (0..8u8).map(|i| (i * 2) | ((i * 2 + 1) << 4)).collect(); - let qh_sub: Vec = vec![0b11_10_01_00, 0b00_01_10_11, 0b10_10_01_01, 0b11_00_11_00]; - let x_sub: Vec = (0..16).map(|i| (i as f32 - 8.0) * 0.125).collect(); - - let s = scalar_subblock_dot_16(&ql_sub, &qh_sub, &x_sub); - let g = q6_subblock_dot_16(&ql_sub, &qh_sub, &x_sub); - // Same arithmetic, possibly different summation order — allow - // small relative drift. - let rel = ((s - g).abs() / s.abs().max(1e-6)) as f64; - assert!(rel < 1e-5, "scalar={s} neon={g}"); - } - - #[test] - fn q6_subblock_zero_input_zero_output() { - let ql = vec![0xFFu8; 8]; - let qh = vec![0xFFu8; 4]; - let x = vec![0.0f32; 16]; - assert_eq!(q6_subblock_dot_16(&ql, &qh, &x), 0.0); - } - - #[test] - fn q6_subblock_zero_weights_zero_output() { - // 6-bit value 32 = (lo4=0, hi2=2). Pack lo=0, hi=2 everywhere - // → all 6-bit values = 32 → after -32 offset all coefficients - // are zero → result zero regardless of x. - let ql = vec![0x00u8; 8]; - let qh = vec![0b10_10_10_10u8; 4]; // every hi2 = 2 - let x: Vec = (0..16).map(|i| i as f32 + 1.0).collect(); - let got = q6_subblock_dot_16(&ql, &qh, &x); - assert!(got.abs() < 1e-5, "expected ~0, got {got}"); - } -} - #[cfg(test)] mod tests { use super::*; @@ -291,6 +112,37 @@ mod tests { ); } + #[test] + fn q6k_round_trip_approximates_f32_matvec() { + // quantize_q6_k (writer) and dispatch (reader) must agree on the + // block layout AND approximate the original f32 matvec — this is + // the pair that broke when the writer packed a private interleaved + // layout while GGUF data arrived planar. + let hidden = 512; + let rows = 8; + let matrix: Vec = (0..rows * hidden) + .map(|i| ((i as f32 * 0.037).sin()) * 0.05) + .collect(); + let x: Vec = (0..hidden).map(|i| ((i as f32 * 0.013).cos()) * 0.5).collect(); + + let q6k = quantize_q6_k(&matrix); + let got = dispatch(&q6k, &x, rows, hidden); + + for r in 0..rows { + let want: f32 = matrix[r * hidden..(r + 1) * hidden] + .iter() + .zip(&x) + .map(|(w, xi)| w * xi) + .sum(); + let tol = 0.02 * want.abs().max(0.5); + assert!( + (got[r] - want).abs() < tol, + "row {r}: got {}, want {want} (tol {tol})", + got[r] + ); + } + } + // ── local f16_to_f32 edge cases ── #[test] diff --git a/crates/larql-inference/src/kv_dispatch/helpers.rs b/crates/larql-inference/src/kv_dispatch/helpers.rs index ca597e518..17afd5e56 100644 --- a/crates/larql-inference/src/kv_dispatch/helpers.rs +++ b/crates/larql-inference/src/kv_dispatch/helpers.rs @@ -24,6 +24,8 @@ use ndarray::Array2; use super::{EngineBackend, KvHandle}; use crate::async_compute_backend::AsyncComputeBackend; use crate::ffn::FfnBackend; +use crate::forward::layer::apply_layer_scalar; +use crate::forward::ple::{apply_per_layer_embedding, precompute_per_layer_inputs}; use crate::forward::{embed_tokens_pub, run_ffn}; /// Per-layer FFN dispatch for the KV-cached engine path, MoE-aware. @@ -70,7 +72,10 @@ pub fn kv_prefill_via_dispatch( return None; } let h = embed_tokens_pub(&weights, prompt_ids); - kv_prefill_from_hidden_via_dispatch(backend, weights, ffn, &h, window, index) + // Per-Layer Embedding inputs for Gemma-4 archs (empty Vec — and a + // per-layer no-op — for everything else). Matches `kv_prefill_run`. + let ple_inputs = precompute_per_layer_inputs(&weights, &h, prompt_ids); + prefill_from_hidden_inner(backend, weights, ffn, &h, window, index, &ple_inputs) } /// Multi-modal-aware peer of [`kv_prefill_via_dispatch`]. Takes @@ -84,6 +89,12 @@ pub fn kv_prefill_via_dispatch( /// produce bit-identical output by construction — the former is a /// two-line wrapper around the latter. Pinned by tests at the bottom /// of this module. +/// PLE caveat: this entry point has no token ids, so Gemma-4 Per-Layer +/// Embeddings cannot be computed here — it passes an empty PLE slab +/// (per-layer no-op). Gemma-4 callers must use the token-id entry +/// point; MM (Gemma-3-only today) is unaffected because Gemma 3 has no +/// PLE. `layer_scalar` IS still applied (weights-driven, no token ids +/// needed). pub fn kv_prefill_from_hidden_via_dispatch( backend: &dyn EngineBackend, weights: larql_models::WeightsView, @@ -91,6 +102,23 @@ pub fn kv_prefill_from_hidden_via_dispatch( initial_hidden: &Array2, window: Option, index: Option<&larql_vindex::VectorIndex>, +) -> Option<(Array2, Vec)> { + prefill_from_hidden_inner(backend, weights, ffn, initial_hidden, window, index, &[]) +} + +/// Shared prefill body. Runs, per layer: attention → FFN → +/// per-layer embedding → layer_scalar — all four steps are required +/// for Gemma-4 correctness (see `larql-compute::forward::layer`), and +/// the last two are no-ops on every other arch. Mirrors +/// `kv_prefill_run` in `larql-kv/src/generation.rs`. +fn prefill_from_hidden_inner( + backend: &dyn EngineBackend, + weights: larql_models::WeightsView, + ffn: &dyn FfnBackend, + initial_hidden: &Array2, + window: Option, + index: Option<&larql_vindex::VectorIndex>, + ple_inputs: &[Array2], ) -> Option<(Array2, Vec)> { if initial_hidden.nrows() == 0 { return None; @@ -114,7 +142,10 @@ pub fn kv_prefill_from_hidden_via_dispatch( } handles.push(handle); - h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let h_ffn = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let mut h_out = apply_per_layer_embedding(&weights, &h_ffn, layer, ple_inputs.get(layer)); + apply_layer_scalar(&weights, &mut h_out, layer); + h = h_out; } Some((last_row_as_2d(&h), handles)) @@ -148,6 +179,10 @@ pub fn kv_decode_step_via_dispatch( "kv_decode_step_via_dispatch: handles.len() must equal weights.num_layers" ); let h_new = embed_tokens_pub(&weights, &[token_id]); + // PLE inputs are per-token — recompute for this single-token step. + // Empty (and a no-op below) on non-Gemma-4 archs. Matches + // `kv_decode_step_run`. + let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[token_id]); let mut h_step = h_new; for (layer, handle) in handles.iter_mut().enumerate().take(num_layers) { @@ -164,7 +199,10 @@ pub fn kv_decode_step_via_dispatch( if let Some(w) = window { backend.clip_kv(handle, w); } - h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let h_ffn = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let mut h_out = apply_per_layer_embedding(&weights, &h_ffn, layer, ple_inputs.get(layer)); + apply_layer_scalar(&weights, &mut h_out, layer); + h_step = h_out; } Some(h_step) @@ -200,7 +238,9 @@ pub fn kv_prefill_via_dispatch_async( return None; } let h = embed_tokens_pub(&weights, prompt_ids); - kv_prefill_from_hidden_via_dispatch_async(backend, weights, ffn, &h, window, index) + // Gemma-4 PLE inputs — same recipe as the sync path. + let ple_inputs = precompute_per_layer_inputs(&weights, &h, prompt_ids); + prefill_from_hidden_inner_async(backend, weights, ffn, &h, window, index, &ple_inputs) } /// Async multi-modal-aware peer of [`kv_prefill_via_dispatch_async`]. @@ -211,6 +251,9 @@ pub fn kv_prefill_via_dispatch_async( /// Bit-identity contract: same as the sync peer. Pinned by the parity /// test at the bottom of this module — sync vs async must agree on /// CPU paths, MM vs text must agree when text is the input. +/// PLE caveat: same as the sync from-hidden peer — no token ids here, +/// so Gemma-4 PLE is skipped (empty slab, per-layer no-op); +/// `layer_scalar` is still applied. pub fn kv_prefill_from_hidden_via_dispatch_async( backend: &dyn AsyncComputeBackend, weights: larql_models::WeightsView, @@ -218,6 +261,20 @@ pub fn kv_prefill_from_hidden_via_dispatch_async( initial_hidden: &Array2, window: Option, index: Option<&larql_vindex::VectorIndex>, +) -> Option<(Array2, Vec)> { + prefill_from_hidden_inner_async(backend, weights, ffn, initial_hidden, window, index, &[]) +} + +/// Shared async prefill body — attention → FFN → per-layer embedding → +/// layer_scalar per layer, mirroring [`prefill_from_hidden_inner`]. +fn prefill_from_hidden_inner_async( + backend: &dyn AsyncComputeBackend, + weights: larql_models::WeightsView, + ffn: &dyn FfnBackend, + initial_hidden: &Array2, + window: Option, + index: Option<&larql_vindex::VectorIndex>, + ple_inputs: &[Array2], ) -> Option<(Array2, Vec)> { if initial_hidden.nrows() == 0 { return None; @@ -242,7 +299,10 @@ pub fn kv_prefill_from_hidden_via_dispatch_async( handles.push(handle); let h_post_attn = backend.read_hidden(h_post_attn_handle); - h = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let h_ffn = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let mut h_out = apply_per_layer_embedding(&weights, &h_ffn, layer, ple_inputs.get(layer)); + apply_layer_scalar(&weights, &mut h_out, layer); + h = h_out; } backend.flush().ok()?; @@ -272,6 +332,8 @@ pub fn kv_decode_step_via_dispatch_async( "kv_decode_step_via_dispatch_async: handles.len() must equal weights.num_layers" ); let h_new = embed_tokens_pub(&weights, &[token_id]); + // Per-token Gemma-4 PLE inputs — same recipe as the sync decode step. + let ple_inputs = precompute_per_layer_inputs(&weights, &h_new, &[token_id]); let mut h_step = h_new; for (layer, handle) in handles.iter_mut().enumerate().take(num_layers) { @@ -287,7 +349,10 @@ pub fn kv_decode_step_via_dispatch_async( backend.clip_kv(handle, w); } let h_post_attn = backend.read_hidden(h_post_attn_handle); - h_step = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let h_ffn = ffn_or_moe_layer(weights, &h_post_attn, layer, ffn); + let mut h_out = apply_per_layer_embedding(&weights, &h_ffn, layer, ple_inputs.get(layer)); + apply_layer_scalar(&weights, &mut h_out, layer); + h_step = h_out; } backend.flush().ok()?; diff --git a/crates/larql-lql/src/executor/query/explain.rs b/crates/larql-lql/src/executor/query/explain.rs index 872f08fbb..5513c4ef4 100644 --- a/crates/larql-lql/src/executor/query/explain.rs +++ b/crates/larql-lql/src/executor/query/explain.rs @@ -11,17 +11,14 @@ impl Session { layers: Option<&Range>, verbose: bool, ) -> Result, LqlError> { - let (path, _config, patched) = self.require_vindex()?; + let (path, config, patched) = self.require_vindex()?; let (embed, embed_scale) = larql_vindex::load_vindex_embeddings(path) .map_err(|e| LqlError::exec("failed to load embeddings", e))?; let tokenizer = larql_vindex::load_vindex_tokenizer(path) .map_err(|e| LqlError::exec("failed to load tokenizer", e))?; - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_vindex_prompt(config, &tokenizer, prompt)?; if token_ids.is_empty() { return Err(LqlError::Execution("empty prompt".into())); diff --git a/crates/larql-lql/src/executor/query/infer.rs b/crates/larql-lql/src/executor/query/infer.rs index 3ac85cf7c..1bf69c856 100644 --- a/crates/larql-lql/src/executor/query/infer.rs +++ b/crates/larql-lql/src/executor/query/infer.rs @@ -47,10 +47,7 @@ impl Session { weights, tokenizer, .. } = &self.backend { - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_dense_prompt(weights, tokenizer, prompt)?; let start = std::time::Instant::now(); let result = larql_inference::predict(weights, tokenizer, &token_ids, top_k); @@ -87,10 +84,7 @@ impl Session { let tokenizer = larql_vindex::load_vindex_tokenizer(path) .map_err(|e| LqlError::exec("failed to load tokenizer", e))?; - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_vindex_prompt(config, &tokenizer, prompt)?; // Shared INFER pipeline — walk FFN (unlimited features) plus KnnStore // side-channel override. Same code path as `PyVindex::infer`; see ADR diff --git a/crates/larql-lql/src/executor/query/infer_trace.rs b/crates/larql-lql/src/executor/query/infer_trace.rs index 158b0ef32..c19279fca 100644 --- a/crates/larql-lql/src/executor/query/infer_trace.rs +++ b/crates/larql-lql/src/executor/query/infer_trace.rs @@ -47,10 +47,7 @@ impl Session { let mut cb = larql_vindex::SilentLoadCallbacks; let tokenizer = larql_vindex::load_vindex_tokenizer(path) .map_err(|e| LqlError::exec("failed to load tokenizer", e))?; - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_vindex_prompt(config, &tokenizer, prompt)?; let token_strs: Vec> = if with_attention { token_ids @@ -193,10 +190,7 @@ impl Session { prompt: &str, top_k: usize, ) -> Result, LqlError> { - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_dense_prompt(weights, tokenizer, prompt)?; let start = std::time::Instant::now(); let result = larql_inference::predict(weights, tokenizer, &token_ids, top_k); diff --git a/crates/larql-lql/src/executor/query/mod.rs b/crates/larql-lql/src/executor/query/mod.rs index 190e3fba6..201c66c63 100644 --- a/crates/larql-lql/src/executor/query/mod.rs +++ b/crates/larql-lql/src/executor/query/mod.rs @@ -1,7 +1,7 @@ //! Query executor: WALK, INFER, SELECT, DESCRIBE, EXPLAIN. //! //! Each verb lives in its own file. Shared helpers (layer-band -//! resolution) live here because both DESCRIBE and EXPLAIN INFER +//! resolution, prompt tokenization) live here because multiple verbs //! consume them. mod describe; @@ -11,6 +11,43 @@ mod infer_trace; mod select; mod walk; +use crate::error::LqlError; + +/// Tokenize an LQL prompt against a vindex, prepending BOS when the +/// architecture requires it but the tokenizer's post-processor doesn't +/// add it (Gemma 4). Every text-prompt query path must route through +/// this (or [`encode_dense_prompt`]) rather than calling +/// `tokenizer.encode` directly — a silently missing BOS is enough to +/// turn gemma-4 prose into token salad. Legacy vindexes without a +/// recorded `model_config` fall back to the tokenizer's own output. +pub(super) fn encode_vindex_prompt( + config: &larql_vindex::VindexConfig, + tokenizer: &larql_inference::tokenizers::Tokenizer, + prompt: &str, +) -> Result, LqlError> { + match larql_vindex::arch_from_vindex_config(config) { + Some(arch) => larql_inference::encode_prompt(tokenizer, arch.as_ref(), prompt) + .map_err(|e| LqlError::exec("tokenize error", e)), + None => { + let encoding = tokenizer + .encode(prompt, true) + .map_err(|e| LqlError::exec("tokenize error", e))?; + Ok(encoding.get_ids().to_vec()) + } + } +} + +/// [`encode_vindex_prompt`] for the dense Weight backend, where the +/// loaded `ModelWeights` already carries the detected architecture. +pub(super) fn encode_dense_prompt( + weights: &larql_inference::ModelWeights, + tokenizer: &larql_inference::tokenizers::Tokenizer, + prompt: &str, +) -> Result, LqlError> { + larql_inference::encode_prompt(tokenizer, weights.arch.as_ref(), prompt) + .map_err(|e| LqlError::exec("tokenize error", e)) +} + /// Resolve the layer-band boundaries from the vindex config, with a /// family-based default and a final whole-range fallback. pub(super) fn resolve_bands(config: &larql_vindex::VindexConfig) -> larql_vindex::LayerBands { diff --git a/crates/larql-lql/src/executor/query/walk.rs b/crates/larql-lql/src/executor/query/walk.rs index 837efe5f2..0a27f7cb6 100644 --- a/crates/larql-lql/src/executor/query/walk.rs +++ b/crates/larql-lql/src/executor/query/walk.rs @@ -13,7 +13,7 @@ impl Session { mode: Option, compare: bool, ) -> Result, LqlError> { - let (path, _config, patched) = self.require_vindex()?; + let (path, config, patched) = self.require_vindex()?; let top_k = top.unwrap_or(10) as usize; let (embed, embed_scale) = larql_vindex::load_vindex_embeddings(path) @@ -21,10 +21,7 @@ impl Session { let tokenizer = larql_vindex::load_vindex_tokenizer(path) .map_err(|e| LqlError::exec("failed to load tokenizer", e))?; - let encoding = tokenizer - .encode(prompt, true) - .map_err(|e| LqlError::exec("tokenize error", e))?; - let token_ids: Vec = encoding.get_ids().to_vec(); + let token_ids = super::encode_vindex_prompt(config, &tokenizer, prompt)?; if token_ids.is_empty() { return Err(LqlError::Execution("empty prompt".into())); diff --git a/crates/larql-models/src/loading/gguf/constants.rs b/crates/larql-models/src/loading/gguf/constants.rs index b955cae24..7617bb1b2 100644 --- a/crates/larql-models/src/loading/gguf/constants.rs +++ b/crates/larql-models/src/loading/gguf/constants.rs @@ -79,4 +79,20 @@ pub(super) const GGUF_TO_HF_KEY_REPLACEMENTS: &[(&str, &str)] = &[ ("output.", "lm_head."), ]; +/// Gemma 2/3/4 layers carry four norms plus QK-norms; the generic table +/// above maps `ffn_norm.` to `post_attention_layernorm.` (correct for the +/// llama two-norm layout, wrong here — gemma's `ffn_norm` is the pre-FFN +/// norm and `post_attention_norm` is the real post-attention one) and has +/// no entries for the rest. Applied BEFORE the generic table so `ffn_norm.` +/// is consumed by the gemma rule first. Gemma 1 keeps the generic path. +pub(super) const GGUF_TO_HF_KEY_REPLACEMENTS_GEMMA: &[(&str, &str)] = &[ + ("attn_q_norm.", "self_attn.q_norm."), + ("attn_k_norm.", "self_attn.k_norm."), + ("post_attention_norm.", "post_attention_layernorm."), + ("ffn_norm.", "pre_feedforward_layernorm."), + ("post_ffw_norm.", "post_feedforward_layernorm."), + // Gemma 4 per-layer output scalar; larql's key has no `.weight` suffix. + ("layer_output_scale.weight", "layer_scalar"), +]; + // Tensor type constants moved to format::quant::ggml diff --git a/crates/larql-models/src/loading/gguf/loader.rs b/crates/larql-models/src/loading/gguf/loader.rs index 7ab7a9757..a1096d3ca 100644 --- a/crates/larql-models/src/loading/gguf/loader.rs +++ b/crates/larql-models/src/loading/gguf/loader.rs @@ -93,8 +93,15 @@ impl GgufFile { let mut vectors = HashMap::new(); let mut raw_bytes: HashMap> = HashMap::new(); + let arch = self + .metadata + .get(GGUF_GENERAL_ARCHITECTURE) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + for info in &self.tensor_infos { - let key = normalize_gguf_key(&info.name); + let key = normalize_gguf_key_for_arch(&info.name, &arch); if skip_key(&key) { continue; } @@ -292,6 +299,81 @@ impl GgufFile { config[HF_VOCAB_SIZE] = serde_json::json!(vocab_size); } + // ── Gemma 4 per-layer attention geometry ──────────────────────────── + // Gemma 4 GGUFs describe heterogeneous attention (sliding layers vs + // global layers) with per-layer arrays and `*_swa` twin keys; the + // flat mapping above collapses them to one number and drops the + // rest. Re-emit the flat HF-style keys the safetensors detect path + // reads (`detect/parser.rs`) so the reconstructed arch can route + // per layer. Without these, gemma-4 models whose global layers have + // no attn_v tensor (12B, 31B: `attention_k_eq_v`) index past the + // end of the V collection at inference time. + if arch == "gemma4" { + // head_count_kv is a per-layer array (e.g. 8 on sliding layers, + // 1 on global layers for the 12B). Layer 0 is always sliding; + // the first differing value is the global-layer count. + if let Some(GgufValue::Array(arr)) = self + .metadata + .get(&format!("{prefix}{GGUF_ATTENTION_HEAD_COUNT_KV}")) + { + let vals: Vec = arr.iter().filter_map(|x| x.as_u32()).collect(); + if let Some(&sliding) = vals.first() { + config[HF_NUM_KEY_VALUE_HEADS] = serde_json::json!(sliding); + if let Some(&global) = vals.iter().find(|&&v| v != sliding) { + config["num_global_key_value_heads"] = serde_json::json!(global); + } + } + } + // key_length is the global-layer head width; key_length_swa the + // sliding-layer width (the base head_dim the arch expects). + let key_len = get_arch_u32(GGUF_ATTENTION_KEY_LENGTH); + let key_len_swa = get_arch_u32("attention.key_length_swa"); + if key_len_swa > 0 { + config[HF_HEAD_DIM] = serde_json::json!(key_len_swa); + } + if key_len > 0 && key_len != key_len_swa { + config["global_head_dim"] = serde_json::json!(key_len); + // Global layers rotate a quarter of their head dims — + // constant across every Gemma 4 HF config (12B/26B/31B); + // the GGUF carries no equivalent key. + config["partial_rotary_factor"] = serde_json::json!(0.25); + } + // Dual RoPE bases: `rope.freq_base` is the global-layer clock + // (already emitted as rope_theta above), `rope.freq_base_swa` + // the sliding-layer one. + if let Some(swa) = get_arch_f64("rope.freq_base_swa") { + config["rope_local_base_freq"] = serde_json::json!(swa); + } + if let Some(sw) = get_arch_u32_opt("attention.sliding_window").filter(|&v| v > 0) { + config["sliding_window"] = serde_json::json!(sw); + } + // Layers with no attn_v tensor reuse K as V (`attention_k_eq_v` + // in the HF config). The GGUF metadata has no flag for this; + // detect it from the tensor inventory, as llama.cpp does. + let n_blocks = get_arch_u32(GGUF_BLOCK_COUNT) as usize; + let n_v = self + .tensor_infos + .iter() + .filter(|t| t.name().ends_with(".attn_v.weight")) + .count(); + if n_blocks > 0 && n_v > 0 && n_v < n_blocks { + config["attention_k_eq_v"] = serde_json::json!(true); + } + // Final-logit softcap: logits = cap * tanh(logits / cap). + // Monotonic (never changes the argmax) but shapes the softmax + // distribution, so probabilities drift without it. + if let Some(cap) = get_arch_f64("final_logit_softcapping") { + config["final_logit_softcapping"] = serde_json::json!(cap); + } + } + + // RMSNorm epsilon — llama.cpp emits it for every RMSNorm family + // under the arch prefix. Absent → detect_from_json falls back to + // its default. + if let Some(eps) = get_arch_f64("attention.layer_norm_rms_epsilon") { + config["rms_norm_eps"] = serde_json::json!(eps); + } + // ── MLA fields (DeepSeek-V2/V3 family, e.g. Kimi K2) ───────────────── // The HF config exposes `q_lora_rank` / `kv_lora_rank` / // `qk_nope_head_dim` / `qk_rope_head_dim` / `v_head_dim`. llama.cpp @@ -509,6 +591,28 @@ pub fn normalize_gguf_key(name: &str) -> String { .fold(name.to_string(), |acc, (from, to)| acc.replace(from, to)) } +/// As [`normalize_gguf_key`], but arch-aware: gemma 2/3/4 GGUFs use a +/// four-norm layer layout (attn_norm / post_attention_norm / ffn_norm / +/// post_ffw_norm) plus QK-norms and, on Gemma 4, a per-layer output +/// scalar. The generic table maps `ffn_norm.` to the llama-style +/// post-attention slot — actively wrong for gemma — and drops the rest +/// on the floor, so gemma-specific replacements run first. Gemma 1 has +/// the llama two-norm layout and stays on the generic path. +pub fn normalize_gguf_key_for_arch(name: &str, arch: &str) -> String { + let gemma_layout = + matches!(arch, "gemma2" | "gemma3") || arch.starts_with("gemma4"); + let name = if gemma_layout { + GGUF_TO_HF_KEY_REPLACEMENTS_GEMMA + .iter() + .fold(name.to_string(), |acc, (from, to)| acc.replace(from, to)) + } else { + name.to_string() + }; + GGUF_TO_HF_KEY_REPLACEMENTS + .iter() + .fold(name, |acc, (from, to)| acc.replace(from, to)) +} + #[cfg(test)] mod tests { use super::super::constants::*; @@ -532,6 +636,50 @@ mod tests { assert_eq!(normalize_gguf_key("output.weight"), "lm_head.weight"); } + #[test] + fn test_normalize_gguf_key_gemma_layout() { + assert_eq!( + normalize_gguf_key_for_arch("blk.0.attn_q_norm.weight", "gemma4"), + "layers.0.self_attn.q_norm.weight" + ); + assert_eq!( + normalize_gguf_key_for_arch("blk.0.attn_k_norm.weight", "gemma4_unified"), + "layers.0.self_attn.k_norm.weight" + ); + assert_eq!( + normalize_gguf_key_for_arch("blk.0.post_attention_norm.weight", "gemma2"), + "layers.0.post_attention_layernorm.weight" + ); + // gemma's ffn_norm is the PRE-feedforward norm... + assert_eq!( + normalize_gguf_key_for_arch("blk.3.ffn_norm.weight", "gemma3"), + "layers.3.pre_feedforward_layernorm.weight" + ); + // ...while llama's ffn_norm keeps the generic post-attention mapping. + assert_eq!( + normalize_gguf_key_for_arch("blk.3.ffn_norm.weight", "llama"), + "layers.3.post_attention_layernorm.weight" + ); + // Gemma 1 is llama-layout too. + assert_eq!( + normalize_gguf_key_for_arch("blk.3.ffn_norm.weight", "gemma"), + "layers.3.post_attention_layernorm.weight" + ); + assert_eq!( + normalize_gguf_key_for_arch("blk.47.post_ffw_norm.weight", "gemma4"), + "layers.47.post_feedforward_layernorm.weight" + ); + assert_eq!( + normalize_gguf_key_for_arch("blk.5.layer_output_scale.weight", "gemma4"), + "layers.5.layer_scalar" + ); + // Mapped projections are untouched by the gemma pre-pass. + assert_eq!( + normalize_gguf_key_for_arch("blk.0.attn_q.weight", "gemma4"), + "layers.0.self_attn.q_proj.weight" + ); + } + #[test] fn test_load_tensors_swaps_gguf_2d_dims_to_rows_cols() { use std::io::{Seek, Write}; diff --git a/crates/larql-models/src/quant/ggml/legacy.rs b/crates/larql-models/src/quant/ggml/legacy.rs index e34ecaa57..89bca711e 100644 --- a/crates/larql-models/src/quant/ggml/legacy.rs +++ b/crates/larql-models/src/quant/ggml/legacy.rs @@ -12,6 +12,8 @@ use crate::quant::half::f16_to_f32; /// Q4_0: block = f16 scale (2B) + 16 bytes of 4-bit quants. 32 elements per block. /// Each 4-bit value is unsigned [0,15], offset by -8 to give signed [-8, 7]. +/// ggml planar nibble layout (`dequantize_row_q4_0`): the low nibbles of +/// qs[0..16] are elements 0..16 and the high nibbles are elements 16..32. pub fn dequantize_q4_0(data: &[u8], n_elements: usize) -> Result, ModelError> { let block_size = 18; let n_blocks = check_block_input("Q4_0", data, n_elements, 32, block_size)?; @@ -24,8 +26,10 @@ pub fn dequantize_q4_0(data: &[u8], n_elements: usize) -> Result, Model for byte in &quants[..16] { let lo = (byte & 0x0F) as i8 - 8; - let hi = ((byte >> 4) & 0x0F) as i8 - 8; out.push(lo as f32 * scale); + } + for byte in &quants[..16] { + let hi = ((byte >> 4) & 0x0F) as i8 - 8; out.push(hi as f32 * scale); } } @@ -34,6 +38,7 @@ pub fn dequantize_q4_0(data: &[u8], n_elements: usize) -> Result, Model /// Q4_1: block = f16 scale + f16 min + 16 bytes of 4-bit quants. /// value = quant * scale + min +/// Same planar nibble layout as Q4_0 (`dequantize_row_q4_1`). pub(super) fn dequantize_q4_1(data: &[u8], n_elements: usize) -> Result, ModelError> { let block_size = 20; let n_blocks = check_block_input("Q4_1", data, n_elements, 32, block_size)?; @@ -47,8 +52,10 @@ pub(super) fn dequantize_q4_1(data: &[u8], n_elements: usize) -> Result for byte in &quants[..16] { let lo = (byte & 0x0F) as f32; - let hi = ((byte >> 4) & 0x0F) as f32; out.push(lo * scale + min); + } + for byte in &quants[..16] { + let hi = ((byte >> 4) & 0x0F) as f32; out.push(hi * scale + min); } } @@ -86,18 +93,18 @@ pub fn dequantize_q5_0(data: &[u8], n_elements: usize) -> Result, Model let high_bits = u32::from_le_bytes([block[2], block[3], block[4], block[5]]); let quants = &block[6..]; + // ggml planar layout (`dequantize_row_q5_0`): element j is the low + // nibble of qs[j] with high bit qh>>j; element j+16 is the high + // nibble of qs[j] with high bit qh>>(j+16). for (j, &byte) in quants[..16].iter().enumerate() { - let lo_lo4 = byte & 0x0F; - let hi_lo4 = (byte >> 4) & 0x0F; - - let lo_hi1 = ((high_bits >> (j * 2)) & 1) as u8; - let hi_hi1 = ((high_bits >> (j * 2 + 1)) & 1) as u8; - - let lo_combined = lo_lo4 | (lo_hi1 << 4); - let hi_combined = hi_lo4 | (hi_hi1 << 4); - - out.push((lo_combined as i32 - 16) as f32 * scale); - out.push((hi_combined as i32 - 16) as f32 * scale); + let lo4 = byte & 0x0F; + let hi1 = ((high_bits >> j) & 1) as u8; + out.push(((lo4 | (hi1 << 4)) as i32 - 16) as f32 * scale); + } + for (j, &byte) in quants[..16].iter().enumerate() { + let lo4 = (byte >> 4) & 0x0F; + let hi1 = ((high_bits >> (j + 16)) & 1) as u8; + out.push(((lo4 | (hi1 << 4)) as i32 - 16) as f32 * scale); } } Ok(out) @@ -117,18 +124,16 @@ pub fn dequantize_q5_1(data: &[u8], n_elements: usize) -> Result, Model let high_bits = u32::from_le_bytes([block[4], block[5], block[6], block[7]]); let quants = &block[8..]; + // Same planar layout as Q5_0 (`dequantize_row_q5_1`). for (j, &byte) in quants[..16].iter().enumerate() { - let lo_lo4 = byte & 0x0F; - let hi_lo4 = (byte >> 4) & 0x0F; - - let lo_hi1 = ((high_bits >> (j * 2)) & 1) as u8; - let hi_hi1 = ((high_bits >> (j * 2 + 1)) & 1) as u8; - - let lo_combined = lo_lo4 | (lo_hi1 << 4); - let hi_combined = hi_lo4 | (hi_hi1 << 4); - - out.push(lo_combined as f32 * scale + min); - out.push(hi_combined as f32 * scale + min); + let lo4 = byte & 0x0F; + let hi1 = ((high_bits >> j) & 1) as u8; + out.push((lo4 | (hi1 << 4)) as f32 * scale + min); + } + for (j, &byte) in quants[..16].iter().enumerate() { + let lo4 = (byte >> 4) & 0x0F; + let hi1 = ((high_bits >> (j + 16)) & 1) as u8; + out.push((lo4 | (hi1 << 4)) as f32 * scale + min); } } Ok(out) diff --git a/crates/larql-models/src/quant/ggml/mod.rs b/crates/larql-models/src/quant/ggml/mod.rs index 2c89a15ab..10fcccd6e 100644 --- a/crates/larql-models/src/quant/ggml/mod.rs +++ b/crates/larql-models/src/quant/ggml/mod.rs @@ -302,13 +302,14 @@ mod tests { #[test] fn q4_0_basic() { - // Scale = 1.0, quants = 0x12 → lo=2-8=-6, hi=1-8=-7 + // Scale = 1.0, quants = 0x12 → lo=2-8=-6 (elements 0..16), + // hi=1-8=-7 (elements 16..32) — ggml planar nibble layout. let mut block = vec![0x00, 0x3C]; // f16 1.0 block.extend_from_slice(&[0x12; 16]); let result = dequantize_q4_0(&block, 32).unwrap(); assert_eq!(result.len(), 32); assert!((result[0] - (-6.0)).abs() < 0.01); - assert!((result[1] - (-7.0)).abs() < 0.01); + assert!((result[16] - (-7.0)).abs() < 0.01); } #[test] @@ -328,8 +329,8 @@ mod tests { let result = dequantize_q4_0(&data, 64).unwrap(); assert_eq!(result.len(), 64); assert!((result[0] - 0.0).abs() < 0.01); // block 0 - assert!((result[32] - 2.0).abs() < 0.01); // block 1: 1*2.0 = 2.0 - assert!((result[33] - (-14.0)).abs() < 0.01); // block 1: -7*2.0 = -14.0 + assert!((result[32] - 2.0).abs() < 0.01); // block 1 lo plane: 1*2.0 = 2.0 + assert!((result[48] - (-14.0)).abs() < 0.01); // block 1 hi plane: -7*2.0 = -14.0 } // ── Q4_1 ── @@ -345,12 +346,13 @@ mod tests { #[test] fn q4_1_with_offset() { - // Scale=2.0, min=-1.0, quants=0x31 → lo=1*2-1=1, hi=3*2-1=5 + // Scale=2.0, min=-1.0, quants=0x31 → lo=1*2-1=1 (elements 0..16), + // hi=3*2-1=5 (elements 16..32) — planar layout. let mut block = vec![0x00, 0x40, 0x00, 0xBC]; // scale=2.0, min=-1.0 block.extend_from_slice(&[0x31; 16]); let result = dequantize_q4_1(&block, 32).unwrap(); assert!((result[0] - 1.0).abs() < 0.01); - assert!((result[1] - 5.0).abs() < 0.01); + assert!((result[16] - 5.0).abs() < 0.01); } // ── Q8_0 ── @@ -448,15 +450,15 @@ mod tests { #[test] fn q5_0_mixed() { // scale=2.0, high_bits=0x00000001 (bit 0 set), quants[0]=0x53 - // element 0: lo4=3, hi1=bit0=1, combined=3|16=19, value=(19-16)*2=6.0 - // element 1: lo4=5, hi1=bit1=0, combined=5, value=(5-16)*2=-22.0 + // element 0 (lo nibble, hi1=bit0=1): combined=3|16=19, (19-16)*2=6.0 + // element 16 (hi nibble, hi1=bit16=0): combined=5, (5-16)*2=-22.0 let mut block = vec![0x00, 0x40]; // f16 2.0 block.extend_from_slice(&0x00000001u32.to_le_bytes()); // high bits block.push(0x53); // quants[0]: lo=3, hi=5 block.extend_from_slice(&[0x00; 15]); // rest zero let result = dequantize_q5_0(&block, 32).unwrap(); assert!((result[0] - 6.0).abs() < 0.01); - assert!((result[1] - (-22.0)).abs() < 0.01); + assert!((result[16] - (-22.0)).abs() < 0.01); } #[test] @@ -561,6 +563,128 @@ mod tests { ); } + // ── Ground truth vs llama.cpp (real GGUF bytes) ── + // + // Bytes are block 0 of the named tensors in + // google/gemma-4-12B-it-qat-q4_0-gguf; expected values produced by + // gguf-py's dequantize (mirrors ggml). These pin larql's decoders to + // the GGUF spec layout, not merely to internal round-trip + // consistency — the interleaved-layout bug they catch survived every + // synthetic test in this file. + + // token_embd.weight block 0 (Q6_K) + const Q6K_GT_BYTES: [u8; 210] = [ + 200, 200, 8, 0, 0, 128, 200, 4, 192, 4, 200, 204, 140, 140, 0, 8, + 133, 1, 6, 138, 11, 139, 134, 11, 6, 10, 0, 5, 10, 10, 6, 128, + 193, 203, 138, 6, 139, 203, 203, 10, 79, 197, 139, 129, 197, 70, 128, 198, + 4, 192, 20, 56, 8, 68, 52, 76, 204, 208, 12, 156, 172, 192, 192, 68, + 118, 65, 127, 22, 198, 223, 106, 65, 69, 150, 198, 69, 251, 193, 219, 122, + 92, 232, 112, 236, 112, 4, 236, 180, 176, 180, 176, 176, 148, 84, 156, 8, + 3, 31, 243, 7, 9, 13, 13, 0, 252, 3, 7, 4, 3, 12, 252, 1, + 198, 68, 12, 4, 129, 7, 64, 128, 193, 70, 196, 6, 192, 15, 70, 4, + 82, 84, 142, 161, 165, 87, 100, 42, 219, 126, 102, 150, 105, 150, 40, 166, + 166, 96, 21, 82, 169, 149, 88, 188, 113, 135, 145, 162, 23, 106, 104, 171, + 213, 32, 214, 133, 217, 107, 187, 105, 230, 164, 84, 106, 116, 148, 229, 146, + 109, 171, 214, 9, 158, 134, 139, 91, 81, 157, 89, 144, 106, 174, 129, 137, + 222, 26, 31, 216, 175, 41, 178, 71, 33, 213, 33, 223, 33, 224, 33, 128, + 120, 129, + ]; + #[rustfmt::skip] + const Q6K_GT_EXPECTED: [f32; 256] = [ + 6.095886230e-03, -1.828765869e-02, 6.095886230e-03, -1.219177246e-02, -1.219177246e-02, 1.219177246e-02, -1.828765869e-02, 3.047943115e-03, + 1.219177246e-02, 3.047943115e-03, 6.095886230e-03, 9.143829346e-03, -3.047943115e-03, 9.143829346e-03, -2.438354492e-02, 6.095886230e-03, + -2.913475037e-03, 1.806354523e-02, 5.826950073e-03, -5.826950073e-03, 2.913475037e-03, 2.913475037e-03, 1.515007019e-02, 1.223659515e-02, + 5.826950073e-03, -1.515007019e-02, 9.323120117e-03, -2.913475037e-03, -1.515007019e-02, -5.826950073e-03, 1.515007019e-02, -9.323120117e-03, + 2.153730392e-02, 3.473758698e-03, -1.806354523e-02, 1.806354523e-02, 3.473758698e-03, 3.473758698e-03, 3.473758698e-03, -6.947517395e-03, + -1.042127609e-02, -1.458978653e-02, 3.473758698e-03, 1.042127609e-02, -3.473758698e-03, 6.947517395e-03, -0.000000000e+00, 6.947517395e-03, + -1.075744629e-02, -2.868652344e-02, -1.075744629e-02, -2.151489258e-02, 7.171630859e-03, -1.075744629e-02, 3.585815430e-03, 2.510070801e-02, + -1.792907715e-02, -1.434326172e-02, -1.792907715e-02, -1.792907715e-02, -3.585815430e-03, 0.000000000e+00, 0.000000000e+00, 3.585815430e-03, + -7.261276245e-03, -7.261276245e-03, -5.809020996e-02, 0.000000000e+00, 0.000000000e+00, -1.452255249e-02, 2.178382874e-02, 0.000000000e+00, + -7.261276245e-03, 2.904510498e-02, 2.178382874e-02, -7.261276245e-03, 1.452255249e-02, -1.452255249e-02, 0.000000000e+00, 0.000000000e+00, + -7.350921631e-03, -0.000000000e+00, 1.470184326e-02, 7.350921631e-03, -0.000000000e+00, 7.350921631e-03, 7.350921631e-03, -1.470184326e-02, + -1.470184326e-02, 2.940368652e-02, 1.470184326e-02, -0.000000000e+00, 1.470184326e-02, -0.000000000e+00, -0.000000000e+00, -7.350921631e-03, + -6.992340088e-03, -6.992340088e-03, 1.398468018e-02, 0.000000000e+00, 1.398468018e-02, -6.992340088e-03, -6.992340088e-03, -5.593872070e-02, + 3.496170044e-02, -6.992340088e-03, -1.398468018e-02, 1.398468018e-02, -6.992340088e-03, 6.992340088e-03, -4.195404053e-02, 2.097702026e-02, + -0.000000000e+00, 6.364822388e-03, 4.932737350e-02, 2.068567276e-02, -0.000000000e+00, -6.364822388e-03, 2.068567276e-02, -6.364822388e-03, + 6.364822388e-03, -2.068567276e-02, -0.000000000e+00, -1.432085037e-02, 3.500652313e-02, 6.364822388e-03, 6.364822388e-03, -6.364822388e-03, + 7.395744324e-03, 2.292680740e-02, -1.109361649e-02, 7.395744324e-03, 7.395744324e-03, -2.292680740e-02, -1.922893524e-02, 1.109361649e-02, + -3.697872162e-03, 1.922893524e-02, 1.922893524e-02, -3.697872162e-03, 1.553106308e-02, 2.292680740e-02, 3.697872162e-03, -7.395744324e-03, + -3.854751587e-03, 2.312850952e-02, 0.000000000e+00, -3.854751587e-03, 0.000000000e+00, 3.854751587e-03, 2.698326111e-02, 1.927375793e-02, + -1.541900635e-02, -1.156425476e-02, -1.541900635e-02, -3.083801270e-02, 3.854751587e-03, 3.854751587e-03, -3.854751587e-03, -7.709503174e-03, + 9.614467621e-03, 1.257276535e-02, 9.614467621e-03, 6.656169891e-03, -6.656169891e-03, -9.614467621e-03, -9.614467621e-03, -0.000000000e+00, + 2.958297729e-03, 9.614467621e-03, 6.656169891e-03, -2.958297729e-03, 9.614467621e-03, 2.958297729e-03, 2.958297729e-03, 2.292680740e-02, + 1.627063751e-02, 2.958297729e-03, -2.958297729e-03, 2.958297729e-03, 1.257276535e-02, -6.656169891e-03, 0.000000000e+00, 0.000000000e+00, + -2.292680740e-02, 1.627063751e-02, 2.958297729e-03, -1.922893524e-02, 0.000000000e+00, 2.292680740e-02, -1.922893524e-02, 2.958297729e-03, + 6.656169891e-03, -2.958297729e-03, 6.656169891e-03, 2.292680740e-02, 2.958297729e-03, -9.614467621e-03, -1.627063751e-02, -2.958297729e-03, + -2.958297729e-03, -6.656169891e-03, 2.958297729e-03, -2.958297729e-03, -2.292680740e-02, 2.958297729e-03, -9.614467621e-03, 6.656169891e-03, + 3.585815430e-03, 1.004028320e-02, -6.454467773e-03, -1.290893555e-02, -6.454467773e-03, -2.294921875e-02, -1.290893555e-02, -3.585815430e-03, + -3.585815430e-03, -3.585815430e-03, -3.585815430e-03, -3.585815430e-03, 6.454467773e-03, 3.585815430e-03, -1.649475098e-02, -2.294921875e-02, + -1.183319092e-02, 2.292680740e-02, -2.292680740e-02, -0.000000000e+00, -1.183319092e-02, 1.183319092e-02, -0.000000000e+00, 1.183319092e-02, + -2.292680740e-02, -0.000000000e+00, 1.183319092e-02, 1.183319092e-02, 1.183319092e-02, -0.000000000e+00, -2.292680740e-02, -0.000000000e+00, + -1.147460938e-02, 1.147460938e-02, 4.589843750e-02, -9.179687500e-02, 2.294921875e-02, 0.000000000e+00, 1.147460938e-02, -2.294921875e-02, + -1.147460938e-02, 1.147460938e-02, -1.147460938e-02, 0.000000000e+00, -1.147460938e-02, 0.000000000e+00, 1.147460938e-02, 0.000000000e+00, + ]; + + // blk.0.attn_q.weight block 0 (Q4_0) + const Q40_GT_BYTES: [u8; 18] = [ + 0, 154, 41, 7, 134, 135, 147, 165, 185, 169, 135, 217, 170, 152, 41, 115, + 39, 137, + ]; + #[rustfmt::skip] + const Q40_GT_EXPECTED: [f32; 32] = [ + -2.929687500e-03, 2.929687500e-03, 5.859375000e-03, 2.929687500e-03, 1.464843750e-02, 8.789062500e-03, -2.929687500e-03, -2.929687500e-03, + 2.929687500e-03, -2.929687500e-03, -5.859375000e-03, -0.000000000e+00, -2.929687500e-03, 1.464843750e-02, 2.929687500e-03, -2.929687500e-03, + 1.757812500e-02, 2.343750000e-02, -0.000000000e+00, -0.000000000e+00, -2.929687500e-03, -5.859375000e-03, -8.789062500e-03, -5.859375000e-03, + -0.000000000e+00, -1.464843750e-02, -5.859375000e-03, -2.929687500e-03, 1.757812500e-02, 2.929687500e-03, 1.757812500e-02, -0.000000000e+00, + ]; + + #[test] + fn q6_k_matches_llama_cpp_ground_truth() { + let got = dequantize_q6_k(&Q6K_GT_BYTES, 256).unwrap(); + for (i, (g, e)) in got.iter().zip(Q6K_GT_EXPECTED.iter()).enumerate() { + assert!( + (g - e).abs() <= 1e-7 + 1e-5 * e.abs(), + "Q6_K element {i}: got {g}, expected {e}" + ); + } + } + + #[test] + fn q4_0_matches_llama_cpp_ground_truth() { + let got = dequantize_q4_0(&Q40_GT_BYTES, 32).unwrap(); + for (i, (g, e)) in got.iter().zip(Q40_GT_EXPECTED.iter()).enumerate() { + assert!( + (g - e).abs() <= 1e-7 + 1e-5 * e.abs(), + "Q4_0 element {i}: got {g}, expected {e}" + ); + } + } + + #[test] + fn q6k_row_dot_matches_ground_truth() { + // dot(decode(block), x) must equal dot(ground_truth, x). + let x: Vec = (0..256).map(|i| ((i as f32) * 0.017).sin()).collect(); + let expected: f32 = Q6K_GT_EXPECTED.iter().zip(&x).map(|(w, xi)| w * xi).sum(); + let got = q6k_row_dot(&Q6K_GT_BYTES, &x).unwrap(); + assert!( + (got - expected).abs() < 1e-4, + "q6k_row_dot: got {got}, expected {expected}" + ); + } + + #[test] + fn q6k_row_scaled_add_matches_ground_truth() { + let mut out = vec![0.0f32; 256]; + q6k_row_scaled_add(&Q6K_GT_BYTES, 2.0, &mut out).unwrap(); + for (i, (g, e)) in out.iter().zip(Q6K_GT_EXPECTED.iter()).enumerate() { + let want = 2.0 * e; + assert!( + (g - want).abs() <= 1e-7 + 1e-5 * want.abs(), + "scaled_add element {i}: got {g}, expected {want}" + ); + } + } + // ── Bounds-check rejection (no panics on malformed input) ── fn assert_short_buffer(res: Result, ModelError>, fmt: &str) { diff --git a/crates/larql-models/src/quant/ggml/q6_k.rs b/crates/larql-models/src/quant/ggml/q6_k.rs index c1f7fc03c..8e22255d4 100644 --- a/crates/larql-models/src/quant/ggml/q6_k.rs +++ b/crates/larql-models/src/quant/ggml/q6_k.rs @@ -7,6 +7,38 @@ use crate::ModelError; use super::check_block_input; use crate::quant::half::f16_to_f32; +/// Decode the 16 signed 6-bit values of scale sub-block `j` (0..16) from +/// one 210-byte Q6_K super-block, following ggml's planar layout +/// (`dequantize_row_q6_K` in llama.cpp): within each 128-element half, +/// the low nibbles of ql[0..32] / ql[32..64] hold elements 0..31 / 32..63 +/// and the high nibbles hold elements 64..95 / 96..127; qh[l] carries the +/// two high bits for elements l, l+32, l+64, l+96 at shifts 0/2/4/6. +/// Scale sub-block j covers elements j*16 .. j*16+16, which always lie in +/// a single nibble plane, so the 16 source bytes are contiguous. +#[inline] +pub fn q6k_subblock_vals(block: &[u8], j: usize) -> [i8; 16] { + let ql = &block[0..128]; + let qh = &block[128..192]; + let half = j / 8; // which 128-element half + let g = j % 8; // scale group within the half + let plane = g / 2; // q1/q2/q3/q4 in ggml's naming + let lbase = (g % 2) * 16; // byte offset within the 32-byte plane row + let ql_off = half * 64 + (plane & 1) * 32 + lbase; + let qh_off = half * 32 + lbase; + let shift = (plane as u32) * 2; + let mut out = [0i8; 16]; + for (i, o) in out.iter_mut().enumerate() { + let lo4 = if plane < 2 { + ql[ql_off + i] & 0x0F + } else { + ql[ql_off + i] >> 4 + }; + let hi2 = (qh[qh_off + i] >> shift) & 0x03; + *o = ((lo4 as i32 | ((hi2 as i32) << 4)) - 32) as i8; + } + out +} + pub fn q6k_row_dot(data: &[u8], x: &[f32]) -> Result { const BLOCK: usize = 210; const SUPER: usize = 256; @@ -40,23 +72,13 @@ pub(super) fn q6k_row_dot_scalar(data: &[u8], x: &[f32], n_blocks: usize) -> f32 let mut acc = 0.0f32; for sb in 0..n_blocks { let block = &data[sb * 210..(sb + 1) * 210]; - let ql = &block[0..128]; - let qh = &block[128..192]; let scales = &block[192..208]; let d = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); for (j, &sc_byte) in scales[..16].iter().enumerate() { let sc = d * (sc_byte as i8) as f32; - for i in 0..16 { - let idx = j * 16 + i; - let lo4 = if idx % 2 == 0 { - ql[idx / 2] & 0x0F - } else { - (ql[idx / 2] >> 4) & 0x0F - }; - let hi2_byte = qh[idx / 4]; - let hi2 = (hi2_byte >> ((idx % 4) * 2)) & 0x03; - let val = ((lo4 as i32) | ((hi2 as i32) << 4)) - 32; - acc += sc * (val as f32) * x[sb * 256 + j * 16 + i]; + let vals = q6k_subblock_vals(block, j); + for (i, &v) in vals.iter().enumerate() { + acc += sc * (v as f32) * x[sb * 256 + j * 16 + i]; } } } @@ -78,36 +100,16 @@ unsafe fn q6k_row_dot_neon(data: &[u8], x: &[f32], n_blocks: usize) -> f32 { let x_ptr = x.as_ptr(); for sb in 0..n_blocks { let block = data.as_ptr().add(sb * BLOCK); - let ql = block; - let qh = block.add(128); let scales = block.add(192); let d = f16_to_f32(u16::from_le_bytes([*block.add(208), *block.add(209)])); let sb_base = x_ptr.add(sb * 256); - // 16 scale subblocks × 16 elements = 256 super-block elements. - // Each subblock j covers ql[j*8..(j+1)*8] (8 bytes → 16 nibbles) and - // qh[j*4..(j+1)*4] (4 bytes → 16 two-bit pairs). + // 16 scale subblocks × 16 elements = 256 super-block elements, + // decoded through the shared planar-layout helper (ggml's + // `dequantize_row_q6_K` ordering), then widened and FMA'd. + let block_slice = std::slice::from_raw_parts(block, 210); for j in 0..16 { let sc = d * (*(scales.add(j) as *const i8)) as f32; - let ql_j = ql.add(j * 8); - let qh_j = qh.add(j * 4); - // Decode 16 signed 6-bit vals via scalar extract → i8 stack array. - // Widening i8 → i32 → f32 then SIMDs. - let mut vals = [0i8; 16]; - for chunk in 0..4 { - let ql_b0 = *ql_j.add(chunk * 2); - let ql_b1 = *ql_j.add(chunk * 2 + 1); - let qh_b = *qh_j.add(chunk); - let base = chunk * 4; - // Even idx: low nibble; odd idx: high nibble. hi2 = (qh >> (k*2)) & 3. - let lo0 = (ql_b0 & 0x0F) as u16 | (((qh_b & 0x03) as u16) << 4); - let lo1 = ((ql_b0 >> 4) & 0x0F) as u16 | ((((qh_b >> 2) & 0x03) as u16) << 4); - let lo2 = (ql_b1 & 0x0F) as u16 | ((((qh_b >> 4) & 0x03) as u16) << 4); - let lo3 = ((ql_b1 >> 4) & 0x0F) as u16 | ((((qh_b >> 6) & 0x03) as u16) << 4); - vals[base] = (lo0 as i16 - 32) as i8; - vals[base + 1] = (lo1 as i16 - 32) as i8; - vals[base + 2] = (lo2 as i16 - 32) as i8; - vals[base + 3] = (lo3 as i16 - 32) as i8; - } + let vals = q6k_subblock_vals(block_slice, j); // Widen i8×16 → i16×8 × 2 → i32×4 × 4 → f32×4 × 4. let vals_i8 = vld1q_s8(vals.as_ptr()); let lo_i16 = vmovl_s8(vget_low_s8(vals_i8)); @@ -155,23 +157,13 @@ pub fn q6k_row_scaled_add(data: &[u8], alpha: f32, out: &mut [f32]) -> Result<() } for sb in 0..n_blocks { let block = &data[sb * block_size..(sb + 1) * block_size]; - let ql = &block[0..128]; - let qh = &block[128..192]; let scales = &block[192..208]; let d = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); for (j, &sc_byte) in scales[..16].iter().enumerate() { let sc = d * (sc_byte as i8) as f32; - for i in 0..16 { - let idx = j * 16 + i; - let lo4 = if idx % 2 == 0 { - ql[idx / 2] & 0x0F - } else { - (ql[idx / 2] >> 4) & 0x0F - }; - let hi2_byte = qh[idx / 4]; - let hi2 = (hi2_byte >> ((idx % 4) * 2)) & 0x03; - let val = ((lo4 as i32) | ((hi2 as i32) << 4)) - 32; - out[sb * 256 + j * 16 + i] += alpha * sc * (val as f32); + let vals = q6k_subblock_vals(block, j); + for (i, &v) in vals.iter().enumerate() { + out[sb * 256 + j * 16 + i] += alpha * sc * (v as f32); } } } @@ -188,24 +180,14 @@ pub fn dequantize_q6_k(data: &[u8], n_elements: usize) -> Result, Model for sb in 0..n_blocks { let block = &data[sb * block_size..(sb + 1) * block_size]; - let ql = &block[0..128]; // lower 4 bits - let qh = &block[128..192]; // upper 2 bits let scales = &block[192..208]; // 16 int8 scales let d = f16_to_f32(u16::from_le_bytes([block[208], block[209]])); for (j, &sc_byte) in scales[..16].iter().enumerate() { let sc = d * (sc_byte as i8) as f32; - for i in 0..16 { - let idx = j * 16 + i; - let lo4 = if idx % 2 == 0 { - ql[idx / 2] & 0x0F - } else { - (ql[idx / 2] >> 4) & 0x0F - }; - let hi2_byte = qh[idx / 4]; - let hi2 = (hi2_byte >> ((idx % 4) * 2)) & 0x03; - let val = ((lo4 as i32) | ((hi2 as i32) << 4)) - 32; - out.push(sc * val as f32); + let vals = q6k_subblock_vals(block, j); + for &v in vals.iter() { + out.push(sc * v as f32); } } } diff --git a/crates/larql-models/src/quant/ggml/quantize.rs b/crates/larql-models/src/quant/ggml/quantize.rs index 0545b9320..d7dd372d0 100644 --- a/crates/larql-models/src/quant/ggml/quantize.rs +++ b/crates/larql-models/src/quant/ggml/quantize.rs @@ -30,10 +30,12 @@ pub fn quantize_q4_0(data: &[f32]) -> Vec { let scale_f16 = crate::quant::half::f32_to_f16(scale); out.extend_from_slice(&scale_f16.to_le_bytes()); - // Quantize: each value → round(val/scale) + 8, clamp to [0, 15] + // Quantize: each value → round(val/scale) + 8, clamp to [0, 15]. + // ggml planar nibble layout (`quantize_row_q4_0_ref`): byte j packs + // element j in its low nibble and element j+16 in its high nibble. for j in 0..16 { - let lo_val = block[j * 2]; - let hi_val = block[j * 2 + 1]; + let lo_val = block[j]; + let hi_val = block[j + 16]; let lo = ((lo_val * inv_scale).round() as i32 + 8).clamp(0, 15) as u8; let hi = ((hi_val * inv_scale).round() as i32 + 8).clamp(0, 15) as u8; out.push(lo | (hi << 4)); diff --git a/crates/larql-vindex/src/format/weights/load/mod.rs b/crates/larql-vindex/src/format/weights/load/mod.rs index 1351e0a47..3d7af208d 100644 --- a/crates/larql-vindex/src/format/weights/load/mod.rs +++ b/crates/larql-vindex/src/format/weights/load/mod.rs @@ -159,6 +159,23 @@ pub fn load_model_weights_kquant_shard( q4k::load_model_weights_kquant_shard(dir, callbacks, expert_filter) } +/// Reconstruct the model architecture from a vindex's recorded config, +/// without loading any weights. Same `build_arch_json` → +/// `detect_from_json` path the weight loaders use, exposed for callers +/// that need arch-level facts before (or instead of) a full weight +/// load — e.g. prompt tokenization, where Gemma 4's BOS token must be +/// prepended manually because the shipped tokenizer.json's +/// post-processor doesn't add it. Returns `None` for legacy vindexes +/// that predate `model_config`. +pub fn arch_from_vindex_config( + config: &crate::VindexConfig, +) -> Option> { + let model_cfg = config.model_config.as_ref()?; + Some(larql_models::detect_from_json(&arch::build_arch_json( + config, model_cfg, + ))) +} + /// Find the tokenizer path near a model or vindex directory. pub fn find_tokenizer_path(dir: &Path) -> Option { let p = dir.join(TOKENIZER_JSON); diff --git a/crates/larql-vindex/src/format/weights/mod.rs b/crates/larql-vindex/src/format/weights/mod.rs index 412994909..f52630a87 100644 --- a/crates/larql-vindex/src/format/weights/mod.rs +++ b/crates/larql-vindex/src/format/weights/mod.rs @@ -27,7 +27,7 @@ pub mod write_layers; pub(crate) use capabilities::ensure_extract_level_supported; pub use load::{ - find_tokenizer_path, load_model_weights, load_model_weights_kquant, + arch_from_vindex_config, find_tokenizer_path, load_model_weights, load_model_weights_kquant, load_model_weights_kquant_shard, load_model_weights_with_opts, LoadWeightsOptions, }; pub use manifest::Q4kManifestEntry; diff --git a/crates/larql-vindex/src/lib.rs b/crates/larql-vindex/src/lib.rs index 5fef8f28f..d645f6778 100644 --- a/crates/larql-vindex/src/lib.rs +++ b/crates/larql-vindex/src/lib.rs @@ -95,7 +95,8 @@ pub use format::huggingface::{ CollectionItem, DownloadProgress, PublishCallbacks, PublishOptions, SilentPublishCallbacks, }; pub use format::weights::{ - load_model_weights, load_model_weights_kquant, load_model_weights_kquant_shard, + arch_from_vindex_config, load_model_weights, load_model_weights_kquant, + load_model_weights_kquant_shard, load_model_weights_with_opts, write_model_weights, write_model_weights_kquant, write_model_weights_kquant_with_opts, write_model_weights_with_opts, DownProjFormat, KquantWriteOptions, LoadWeightsOptions, StreamingWeights, WeightSource, WriteWeightsOptions,