ember_tiled Namespace Reference

Functions
__host__ __device__ int	tiled_chunk (int L, int min_chunk, int max_chunks)
__host__ __device__ int	tiled_nchunks (int L, int chunk)
__host__ __device__ int	tiled_grid_chunks (int max_len, int min_chunk, int max_chunks)
__host__ bool	tiled_shape_ok (int heads, int head_dim, int max_len, int heads_max, int min_chunk, int max_chunks)
__host__ size_t	tiled_shm_bytes (int head_dim, int max_len, int min_chunk, int max_chunks)
__device__ __forceinline__ void	tiled_split_block (const float *__restrict__ Qh, const float *__restrict__ K, const float *__restrict__ V, int L, int c, int kvh, int num_kv_heads, int head_dim, float scale, int min_chunk, int max_chunks, float *__restrict__ part)
__device__ __forceinline__ void	tiled_merge_block (const float *__restrict__ base, float *__restrict__ out_h, int L, int head_dim, int min_chunk, int max_chunks)

Variables
constexpr int	kTiledThreads = 128
constexpr int	kTiledMinChunk = 32
constexpr int	kTiledMaxChunks = 64
constexpr int	kTiledChunkMax = 2048
constexpr int	kTiledHeadDimMax = 256

Function Documentation

tiled_chunk()

__host__ __device__ int ember_tiled::tiled_chunk ( int  L, int  min_chunk, int  max_chunks  )

inline

Definition at line 35 of file decode_attn_tiled.cuh.

                                                                                 {
  const int c = (L + max_chunks - 1) / max_chunks;
  return c < min_chunk ? min_chunk : c;
}

References c, and j.

Referenced by tiled_merge_block(), tiled_shape_ok(), tiled_shm_bytes(), and tiled_split_block().

tiled_grid_chunks()

__host__ __device__ int ember_tiled::tiled_grid_chunks ( int  max_len, int  min_chunk, int  max_chunks  )

inline

Definition at line 41 of file decode_attn_tiled.cuh.

                                                                                             {
  const int n = (max_len + min_chunk - 1) / min_chunk;
  return n < max_chunks ? n : max_chunks;
}

References j, and n.

Referenced by ember_kvq::partial_floats().

tiled_merge_block()

__device__ __forceinline__ void ember_tiled::tiled_merge_block	(	const float *__restrict__	base,
		float *__restrict__	out_h,
		int	L,
		int	head_dim,
		int	min_chunk,
		int	max_chunks
	)

Definition at line 152 of file decode_attn_tiled.cuh.

                                                                                               {
  const int nc = tiled_nchunks(L, tiled_chunk(L, min_chunk, max_chunks));  // <= max_chunks <= 64
  const int stride = head_dim + 2;
  __shared__ float w[kTiledMaxChunks];
  __shared__ float inv_s;
  const int tid = threadIdx.x;
  if (tid < 32) {
    const int c0 = tid, c1 = tid + 32;
    const float m0 = c0 < nc ? base[(size_t)c0 * stride] : -INFINITY;
    const float m1 = c1 < nc ? base[(size_t)c1 * stride] : -INFINITY;
    float M = fmaxf(m0, m1);
    for (int o = 16; o > 0; o >>= 1) M = fmaxf(M, __shfl_xor_sync(0xffffffff, M, o));
    const float w0 = c0 < nc ? expf(m0 - M) : 0.0f;
    const float w1 = c1 < nc ? expf(m1 - M) : 0.0f;
    w[c0] = w0;
    w[c1] = w1;
    // Explicit roundings: no contraction into an fma, so every kernel that inlines this adds the same values.
    const float l0 = c0 < nc ? __fmul_rn(w0, base[(size_t)c0 * stride + 1]) : 0.0f;
    const float l1 = c1 < nc ? __fmul_rn(w1, base[(size_t)c1 * stride + 1]) : 0.0f;
    float l = __fadd_rn(l0, l1);
    for (int o = 16; o > 0; o >>= 1) l += __shfl_xor_sync(0xffffffff, l, o);
    if (tid == 0) inv_s = l > 0.0f ? 1.0f / l : 0.0f;
  }
  __syncthreads();
  const float inv = inv_s;
  for (int d = tid; d < head_dim; d += blockDim.x) {
    float acc = 0.0f;
    const float* p = base + 2 + d;
#pragma unroll 4
    for (int c = 0; c < nc; ++c) acc = fmaf(w[c], p[(size_t)c * stride], acc);
    out_h[d] = acc * inv;
  }
}

References c, j, kTiledMaxChunks, p, tiled_chunk(), and tiled_nchunks().

tiled_nchunks()

__host__ __device__ int ember_tiled::tiled_nchunks ( int  L, int  chunk  )

inline

Definition at line 39 of file decode_attn_tiled.cuh.

{ return (L + chunk - 1) / chunk; }

References j.

Referenced by tiled_merge_block().

tiled_shape_ok()

__host__ bool ember_tiled::tiled_shape_ok ( int  heads, int  head_dim, int  max_len, int  heads_max, int  min_chunk, int  max_chunks  )

inline

Definition at line 46 of file decode_attn_tiled.cuh.

                                                    {
  return heads > 0 && heads <= heads_max && head_dim > 0 && head_dim % 4 == 0 && head_dim <= kTiledHeadDimMax &&
         max_len > 0 && max_chunks <= kTiledMaxChunks && tiled_chunk(max_len, min_chunk, max_chunks) <= kTiledChunkMax;
}

References j, kTiledChunkMax, kTiledHeadDimMax, kTiledMaxChunks, and tiled_chunk().

Referenced by ember_kvq::decode_cuda().

tiled_shm_bytes()

__host__ size_t ember_tiled::tiled_shm_bytes ( int  head_dim, int  max_len, int  min_chunk, int  max_chunks  )

inline

Definition at line 52 of file decode_attn_tiled.cuh.

                                                                                                 {
  return (size_t)(head_dim + tiled_chunk(max_len, min_chunk, max_chunks)) * sizeof(float);
}

References j, and tiled_chunk().

tiled_split_block()

__device__ __forceinline__ void ember_tiled::tiled_split_block	(	const float *__restrict__	Qh,
		const float *__restrict__	K,
		const float *__restrict__	V,
		int	L,
		int	c,
		int	kvh,
		int	num_kv_heads,
		int	head_dim,
		float	scale,
		int	min_chunk,
		int	max_chunks,
		float *__restrict__	part
	)

Definition at line 58 of file decode_attn_tiled.cuh.

                                                                                            {
  const int chunk = tiled_chunk(L, min_chunk, max_chunks);
  const int k0 = c * chunk;
  if (k0 >= L) return;  // past nchunks(L): the merge does not read this slot
  const int k1 = min(L, k0 + chunk);
  const int n = k1 - k0;
  const int d4n = head_dim >> 2;
  extern __shared__ float4 tiled_sm4[];
  float4* q4 = tiled_sm4;                       // [head_dim / 4]
  float* s = reinterpret_cast<float*>(q4 + d4n);  // [chunk]
  __shared__ float red[kTiledThreads / 32];
  __shared__ float4 pacc4[kTiledThreads];       // [G][head_dim / 4]
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
  const float4* Q4 = reinterpret_cast<const float4*>(Qh);
  for (int j = tid; j < d4n; j += kTiledThreads) q4[j] = Q4[j];
  __syncthreads();
  const size_t row4 = (size_t)num_kv_heads * d4n;  // float4s per cache position
  const float4* K4 = reinterpret_cast<const float4*>(K) + (size_t)kvh * d4n;
  const float4* V4 = reinterpret_cast<const float4*>(V) + (size_t)kvh * d4n;
  // Scores: quad (tid >> 2) takes key kb + quad; lane (tid & 3) the float4 slices ql, ql + 4, ...
  const int quad = tid >> 2, ql = tid & 3;
  for (int kb = k0; kb < k1; kb += kTiledThreads / 4) {
    const int k = kb + quad;
    float dot = 0.0f;
    if (k < k1) {
      const float4* kr = K4 + (size_t)k * row4;
#pragma unroll 4
      for (int j = ql; j < d4n; j += 4) {
        const float4 a = q4[j];
        const float4 b = __ldg(kr + j);
        dot = fmaf(a.x, b.x, dot);
        dot = fmaf(a.y, b.y, dot);
        dot = fmaf(a.z, b.z, dot);
        dot = fmaf(a.w, b.w, dot);
      }
    }
    dot += __shfl_xor_sync(0xffffffff, dot, 1);
    dot += __shfl_xor_sync(0xffffffff, dot, 2);
    if (k < k1 && ql == 0) s[k - k0] = dot * scale;
  }
  __syncthreads();
  float mx = -INFINITY;
  for (int i = tid; i < n; i += kTiledThreads) mx = fmaxf(mx, s[i]);
  for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffff, mx, o));
  if (lane == 0) red[warp] = mx;
  __syncthreads();
  mx = red[0];
  for (int w = 1; w < kTiledThreads / 32; ++w) mx = fmaxf(mx, red[w]);
  __syncthreads();
  float sum = 0.0f;
  for (int i = tid; i < n; i += kTiledThreads) {
    const float e = expf(s[i] - mx);
    s[i] = e;
    sum += e;
  }
  for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, o);
  if (lane == 0) red[warp] = sum;
  __syncthreads();  // also publishes s[] = e^{s - max} to the value pass
  if (tid == 0) {
    float l = 0.0f;
    for (int w = 0; w < kTiledThreads / 32; ++w) l += red[w];
    part[0] = mx;
    part[1] = l;
  }
  // Values: group g = tid / d4n adds keys g, g + G, ... (G = 128 / d4n groups; threads past G * d4n idle).
  const int G = kTiledThreads / d4n;
  const int g = tid / d4n, j = tid - g * d4n;
  if (g < G) {
    float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
    const float4* vc = V4 + (size_t)k0 * row4 + j;
#pragma unroll 4
    for (int i = g; i < n; i += G) {
      const float w = s[i];
      const float4 v = __ldg(vc + (size_t)i * row4);
      acc.x = fmaf(w, v.x, acc.x);
      acc.y = fmaf(w, v.y, acc.y);
      acc.z = fmaf(w, v.z, acc.z);
      acc.w = fmaf(w, v.w, acc.w);
    }
    pacc4[g * d4n + j] = acc;
  }
  __syncthreads();
  const float* pacc = reinterpret_cast<const float*>(pacc4);
  for (int d = tid; d < head_dim; d += kTiledThreads) {
    float a = 0.0f;
    for (int gg = 0; gg < G; ++gg) a += pacc[gg * head_dim + d];
    part[2 + d] = a;
  }
}

References a, c, j, k, kTiledThreads, n, scale, tiled_chunk(), and v.

Variable Documentation

kTiledChunkMax

constexpr int ember_tiled::kTiledChunkMax = 2048

constexpr

Definition at line 32 of file decode_attn_tiled.cuh.

Referenced by tiled_shape_ok().

kTiledHeadDimMax

constexpr int ember_tiled::kTiledHeadDimMax = 256

constexpr

Definition at line 33 of file decode_attn_tiled.cuh.

Referenced by tiled_shape_ok().

kTiledMaxChunks

constexpr int ember_tiled::kTiledMaxChunks = 64

constexpr

Definition at line 31 of file decode_attn_tiled.cuh.

Referenced by decode_attention_tiles(), tiled_merge_block(), and tiled_shape_ok().

kTiledMinChunk

constexpr int ember_tiled::kTiledMinChunk = 32

constexpr

Definition at line 30 of file decode_attn_tiled.cuh.

Referenced by decode_attention_tiles().

kTiledThreads

constexpr int ember_tiled::kTiledThreads = 128

constexpr

Definition at line 29 of file decode_attn_tiled.cuh.

Referenced by tiled_split_block().