Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 15 additions & 13 deletions crates/larql-compute/csrc/q4_dot.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
}
}
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
}
Expand Down
62 changes: 31 additions & 31 deletions crates/larql-compute/src/cpu/ops/q4_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,11 @@ pub fn quantize_q4_0(data: &[f32]) -> Vec<u8> {
}
};
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));
}
}
Expand Down Expand Up @@ -303,21 +305,29 @@ pub fn quantize_q6_k(data: &[f32]) -> Vec<u8> {
}
}

// 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
Expand Down Expand Up @@ -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;
}
}
}
Expand Down Expand Up @@ -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)
Expand Down
53 changes: 26 additions & 27 deletions crates/larql-compute/src/cpu/ops/q4k_q8k_dot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`:
Expand Down Expand Up @@ -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];
Expand All @@ -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;
}
Expand All @@ -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],
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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);
}

Expand Down
Loading