ember_kvq Namespace Reference

Enumerations
enum	Format : int { kF32 = 0 , kF16 = 1 , kQ8 = 2 , kQ4 = 3 }

Functions
void	write_cuda (int f, void *Kd, void *Ks, void *Vd, void *Vs, const float *k, const float *v, int T, int row, int start_pos, cudaStream_t st, const int *dpos)
size_t	partial_floats (int hq, int hd, int max_len)
void	decode_cuda (int f, const float *q, const void *Kd, const void *Ks, const void *Vd, const void *Vs, float *out, int hq, int hkv, int hd, float scale, int max_len, int seq_len, float *part, cudaStream_t st, const int *dpos)
void	dequant_rows_cuda (int f, const void *Kd, const void *Ks, const void *Vd, const void *Vs, float *K, float *V, int n, int row, cudaStream_t st)
Format	parse_format (const char *s)
const char *	format_name (int f)
size_t	data_bytes (int f, size_t row)
size_t	scale_bytes (int f, size_t row)
std::string	shape_error (int f, int head_dim)
uint16_t	f32_to_f16 (float x)
float	f16_to_f32 (uint16_t h)
void	roundtrip_block (int f, float *x)
void	roundtrip_rows (int f, float *x, size_t rows, size_t row)

Variables
constexpr int	kBlock = 32

Enumeration Type Documentation

Format

enum ember_kvq::Format : int

Enumerator
kF32
kF16
kQ8
kQ4

Definition at line 33 of file kv_quant.h.

: int { kF32 = 0, kF16 = 1, kQ8 = 2, kQ4 = 3 };

Function Documentation

data_bytes()

size_t ember_kvq::data_bytes ( int  f, size_t  row  )

inline

Definition at line 48 of file kv_quant.h.

                                            {
  return f == kF16 ? row * 2 : f == kQ8 ? row : f == kQ4 ? row / 2 : row * 4;
}

References j, kF16, kQ4, and kQ8.

decode_cuda()

void ember_kvq::decode_cuda	(	int	f,
		const float *	q,
		const void *	Kd,
		const void *	Ks,
		const void *	Vd,
		const void *	Vs,
		float *	out,
		int	hq,
		int	hkv,
		int	hd,
		float	scale,
		int	max_len,
		int	seq_len,
		float *	part,
		cudaStream_t	st,
		const int *	dpos
	)

Definition at line 288 of file kv_quant.cu.

                                  {
  if (!ember_tiled::tiled_shape_ok(hq, hd, max_len, 1 << 20, kMin, kMax) || hkv <= 0 || hq % hkv != 0)
    throw std::runtime_error("kv_quant decode: shape not supported (heads " + std::to_string(hq) + "/" +
                             std::to_string(hkv) + ", head_dim " + std::to_string(hd) + ", max_len " +
                             std::to_string(max_len) + ")");
  switch (f) {
    case kF16: decode_launch<kF16>(q, Kd, Ks, Vd, Vs, out, hq, hkv, hd, scale, max_len, seq_len, part, st, dpos); break;
    case kQ8: decode_launch<kQ8>(q, Kd, Ks, Vd, Vs, out, hq, hkv, hd, scale, max_len, seq_len, part, st, dpos); break;
    case kQ4: decode_launch<kQ4>(q, Kd, Ks, Vd, Vs, out, hq, hkv, hd, scale, max_len, seq_len, part, st, dpos); break;
    default: throw std::runtime_error("kv_quant decode: format " + std::to_string(f));
  }
}

References j, kF16, kQ4, kQ8, out, scale, and ember_tiled::tiled_shape_ok().

dequant_rows_cuda()

void ember_kvq::dequant_rows_cuda	(	int	f,
		const void *	Kd,
		const void *	Ks,
		const void *	Vd,
		const void *	Vs,
		float *	K,
		float *	V,
		int	n,
		int	row,
		cudaStream_t	st
	)

Definition at line 303 of file kv_quant.cu.

                                                 {
  if (n <= 0) return;
  switch (f) {
    case kF16: dequant_launch<kF16>(Kd, Ks, Vd, Vs, K, V, n, row, st); break;
    case kQ8: dequant_launch<kQ8>(Kd, Ks, Vd, Vs, K, V, n, row, st); break;
    case kQ4: dequant_launch<kQ4>(Kd, Ks, Vd, Vs, K, V, n, row, st); break;
    default: throw std::runtime_error("kv_quant dequant: format " + std::to_string(f));
  }
}

References j, kF16, kQ4, kQ8, and n.

f16_to_f32()

float ember_kvq::f16_to_f32 ( uint16_t  h )

inline

Definition at line 88 of file kv_quant.h.

                                    {
  const uint32_t sign = static_cast<uint32_t>(h & 0x8000u) << 16;
  const uint32_t e = (h >> 10) & 0x1fu, m = h & 0x3ffu;
  uint32_t u;
  if (e == 0) {
    if (m == 0) {
      u = sign;
    } else {  // subnormal: m * 2^-24, exact in float
      const float v = static_cast<float>(m) * 5.9604644775390625e-8f;
      std::memcpy(&u, &v, 4);
      u |= sign;
    }
  } else if (e == 31) {
    u = sign | 0x7f800000u | (m << 13);
  } else {
    u = sign | ((e + 112u) << 23) | (m << 13);
  }
  float f;
  std::memcpy(&f, &u, 4);
  return f;
}

References j, and v.

Referenced by roundtrip_block(), and roundtrip_rows().

f32_to_f16()

uint16_t ember_kvq::f32_to_f16 ( float  x )

inline

Definition at line 64 of file kv_quant.h.

                                    {
  uint32_t u;
  std::memcpy(&u, &x, 4);
  const uint32_t sign = (u >> 16) & 0x8000u;
  const uint32_t ax = u & 0x7fffffffu;
  if (ax >= 0x7f800000u) return static_cast<uint16_t>(sign | (ax > 0x7f800000u ? 0x7e00u : 0x7c00u));
  if (ax >= 0x477ff000u) return static_cast<uint16_t>(sign | 0x7c00u);  // rounds past 65504: infinity
  if (ax < 0x38800000u) {  // below 2^-14: a half subnormal (or zero)
    if (ax < 0x33000000u) return static_cast<uint16_t>(sign);  // below half the smallest subnormal: 0
    const uint32_t e = ax >> 23;
    const uint32_t m = (ax & 0x7fffffu) | 0x800000u;
    const uint32_t shift = 126 - e;  // r = m * 2^(e - 126) units of 2^-24, m with its implicit bit
    const uint32_t base = m >> shift;
    const uint32_t rem = m & ((1u << shift) - 1u);
    const uint32_t half_ulp = 1u << (shift - 1);
    uint32_t r = base;
    if (rem > half_ulp || (rem == half_ulp && (base & 1u))) ++r;
    return static_cast<uint16_t>(sign | r);
  }
  uint32_t r = ((ax >> 13) - (112u << 10));  // rebias the exponent, keep 10 mantissa bits
  const uint32_t rem = ax & 0x1fffu;
  if (rem > 0x1000u || (rem == 0x1000u && (r & 1u))) ++r;  // a carry into the exponent is correct
  return static_cast<uint16_t>(sign | r);
}

References j.

Referenced by roundtrip_block(), and roundtrip_rows().

format_name()

const char * ember_kvq::format_name ( int  f )

inline

Definition at line 45 of file kv_quant.h.

{ return f == kF16 ? "f16" : f == kQ8 ? "q8" : f == kQ4 ? "q4" : "f32"; }

References j, kF16, kQ4, and kQ8.

Referenced by ember_ir::Interpreter::init(), and ember_ir::Interpreter::kv_format().

parse_format()

Format ember_kvq::parse_format ( const char *  s )

inline

Definition at line 37 of file kv_quant.h.

                                          {
  const std::string v = s ? s : "";
  if (v.empty() || v == "f32") return kF32;
  if (v == "f16") return kF16;
  if (v == "q8") return kQ8;
  if (v == "q4") return kQ4;
  throw std::runtime_error("EMBER_KV=" + v + ": expected f32, f16, q8 or q4");
}

References j, kF16, kF32, kQ4, kQ8, and v.

Referenced by ember_ir::Interpreter::init().

partial_floats()

size_t ember_kvq::partial_floats	(	int	hq,
		int	hd,
		int	max_len
	)

Definition at line 284 of file kv_quant.cu.

                                                   {
  return static_cast<size_t>(hq) * ember_tiled::tiled_grid_chunks(max_len, kMin, kMax) * (hd + 2);
}

References j, and ember_tiled::tiled_grid_chunks().

roundtrip_block()

void ember_kvq::roundtrip_block ( int  f, float *  x  )

inline

Definition at line 111 of file kv_quant.h.

                                             {
  if (f == kQ8) {
    float amax = 0.0f;
    for (int i = 0; i < kBlock; ++i) amax = std::fmax(amax, std::fabs(x[i]));
    const float d = f16_to_f32(f32_to_f16(amax / 127.0f));
    for (int i = 0; i < kBlock; ++i) {
      float q = d != 0.0f ? std::nearbyint(x[i] / d) : 0.0f;
      q = std::fmin(127.0f, std::fmax(-127.0f, q));
      x[i] = q * d;
    }
  } else if (f == kQ4) {
    float mn = x[0], mx = x[0];
    for (int i = 1; i < kBlock; ++i) {
      mn = std::fmin(mn, x[i]);
      mx = std::fmax(mx, x[i]);
    }
    const float d = f16_to_f32(f32_to_f16((mx - mn) / 15.0f));
    const float m = f16_to_f32(f32_to_f16(mn));
    for (int i = 0; i < kBlock; ++i) {
      float q = d != 0.0f ? std::nearbyint((x[i] - m) / d) : 0.0f;
      q = std::fmin(15.0f, std::fmax(0.0f, q));
      x[i] = std::fma(q, d, m);
    }
  }
}

References f16_to_f32(), f32_to_f16(), j, kBlock, kQ4, and kQ8.

Referenced by roundtrip_rows().

roundtrip_rows()

void ember_kvq::roundtrip_rows ( int  f, float *  x, size_t  rows, size_t  row  )

inline

Definition at line 138 of file kv_quant.h.

                                                                     {
  if (f == kF32) return;
  const size_t n = rows * row;
  if (f == kF16) {
    for (size_t i = 0; i < n; ++i) x[i] = f16_to_f32(f32_to_f16(x[i]));
    return;
  }
  for (size_t b = 0; b + kBlock <= n; b += kBlock) roundtrip_block(f, x + b);
}

References f16_to_f32(), f32_to_f16(), j, kBlock, kF16, kF32, n, and roundtrip_block().

Referenced by ember_ir::Interpreter::run_nodes().

scale_bytes()

size_t ember_kvq::scale_bytes ( int  f, size_t  row  )

inline

Definition at line 51 of file kv_quant.h.

                                             {
  return f == kQ8 ? row / kBlock * 2 : f == kQ4 ? row / kBlock * 4 : 0;
}

References j, kBlock, kQ4, and kQ8.

shape_error()

std::string ember_kvq::shape_error ( int  f, int  head_dim  )

inline

Definition at line 55 of file kv_quant.h.

                                                  {
  if (f == kF32) return "";
  if (head_dim % 4 != 0) return "head_dim % 4 != 0";
  if ((f == kQ8 || f == kQ4) && head_dim % kBlock != 0) return "head_dim % 32 != 0 (q8 and q4 blocks sit inside a head)";
  return "";
}

References j, kBlock, kF32, kQ4, and kQ8.

Referenced by ember_ir::Interpreter::init().

write_cuda()

void ember_kvq::write_cuda	(	int	f,
		void *	Kd,
		void *	Ks,
		void *	Vd,
		void *	Vs,
		const float *	k,
		const float *	v,
		int	T,
		int	row,
		int	start_pos,
		cudaStream_t	st,
		const int *	dpos
	)

Definition at line 260 of file kv_quant.cu.

                                                                 {
  if (T <= 0) return;
  if (f == kF16) {
    const size_t total = 2 * static_cast<size_t>(T) * row;
    write_f16_kernel<<<static_cast<unsigned>((total + 255) / 256), 256, 0, st>>>(
        static_cast<__half*>(Kd), static_cast<__half*>(Vd), k, v, T, row, start_pos, dpos);
  } else {
    const size_t threads = 2 * static_cast<size_t>(T) * (row / kBlock) * 32;
    const unsigned grid = static_cast<unsigned>((threads + 127) / 128);
    auto* kd = static_cast<uint8_t*>(Kd);
    auto* ks = static_cast<uint8_t*>(Ks);
    auto* vd = static_cast<uint8_t*>(Vd);
    auto* vs = static_cast<uint8_t*>(Vs);
    if (f == kQ8)
      write_block_kernel<kQ8><<<grid, 128, 0, st>>>(kd, ks, vd, vs, k, v, T, row, start_pos, dpos);
    else if (f == kQ4)
      write_block_kernel<kQ4><<<grid, 128, 0, st>>>(kd, ks, vd, vs, k, v, T, row, start_pos, dpos);
    else
      throw std::runtime_error("kv_quant write: format " + std::to_string(f));
  }
  check(cudaGetLastError(), "write");
}

References j, k, kBlock, kF16, kQ4, kQ8, and v.

Variable Documentation

kBlock

constexpr int ember_kvq::kBlock = 32

constexpr

Definition at line 34 of file kv_quant.h.

Referenced by roundtrip_block(), roundtrip_rows(), scale_bytes(), shape_error(), and write_cuda().