Reese Levine
e8c54893f2
* Start work on flash_attn refactor * Refactor * Split k/v quantization * Refactor and abstract quantization logic for flash_attn and mul_mat * Add quantization support to tile path * formatting * Move to functions, add a check
22 lines
780 B
Cheetah
22 lines
780 B
Cheetah
#ifdef U32_DEQUANT_HELPERS
|
|
fn dequant_q4_0_packed_to_shmem(q_packed: u32, d: f16, dst_idx: u32) {
|
|
let scale = QUANT_OUT_TYPE(d);
|
|
for (var k = 0u; k < 4u; k++) {
|
|
let q_byte = get_byte(q_packed, k);
|
|
let q_hi = (QUANT_OUT_TYPE((q_byte >> 4) & 0xFu) - QUANT_OUT_TYPE(8.0)) * scale;
|
|
let q_lo = (QUANT_OUT_TYPE(q_byte & 0xFu) - QUANT_OUT_TYPE(8.0)) * scale;
|
|
QUANT_SHMEM[dst_idx + k] = q_lo;
|
|
QUANT_SHMEM[dst_idx + k + 16u] = q_hi;
|
|
}
|
|
}
|
|
|
|
fn dequant_q8_0_packed_to_shmem(q_packed: u32, d: f16, dst_idx: u32) {
|
|
let scale = QUANT_OUT_TYPE(d);
|
|
for (var k = 0u; k < 4u; k++) {
|
|
let q_byte = get_byte_i32(q_packed, k);
|
|
let q_val = QUANT_OUT_TYPE(q_byte) * scale;
|
|
QUANT_SHMEM[dst_idx + k] = q_val;
|
|
}
|
|
}
|
|
#endif
|