CUDA KV Cache：分页存储与长上下文 Decode 极致优化

1
template <int HeadDim, int PageSize>
2
__global__ void append_paged_kv(
3
    const half* __restrict__ new_kv,       // [B, 2, Hkv, D]
4
    half* __restrict__ cache,              // [P, 2, T, Hkv, D]
5
    const int* __restrict__ block_table,
6
    const int* __restrict__ sequence_len,
7
    int max_pages,
8
    int num_kv_heads) {
9
  const int batch = blockIdx.x;
10
  const int kind_head = blockIdx.y;
11
  const int kind = kind_head / num_kv_heads;
12
  const int head = kind_head % num_kv_heads;
13
  const int token = sequence_len[batch] - 1;
14
  const int logical_page = token / PageSize;
15
  const int page_offset = token % PageSize;
16
  const int physical_page = block_table[batch * max_pages + logical_page];
17

18
  constexpr int kHalfPerVector = sizeof(uint4) / sizeof(half);
19
  const int vector_index = threadIdx.x;
20
  if (vector_index < HeadDim / kHalfPerVector) {
21
    const size_t source =
22
        (((batch * 2 + kind) * num_kv_heads + head) * HeadDim);
23
    const size_t target =
24
        ((((static_cast<size_t>(physical_page) * 2 + kind) * PageSize
25
            + page_offset) * num_kv_heads + head) * HeadDim);
26
    reinterpret_cast<uint4*>(cache + target)[vector_index] =
27
        reinterpret_cast<const uint4*>(new_kv + source)[vector_index];
28
  }
29
}

1
template <int HeadDim, int PageSize>
2
__global__ void paged_decode(
3
    const half* __restrict__ q,
4
    const half* __restrict__ cache,
5
    const int* __restrict__ page_table,
6
    const int* __restrict__ lengths,
7
    half* __restrict__ out,
8
    int query_heads,
9
    int kv_heads,
10
    int max_pages) {
11
  const int lane = threadIdx.x & 31;
12
  const int warp = threadIdx.x >> 5;
13
  const int query_head = blockIdx.y * (blockDim.x / 32) + warp;
14
  const int batch = blockIdx.x;
15
  if (query_head >= query_heads) return;
16

17
  const int group_size = query_heads / kv_heads;
18
  const int kv_head = query_head / group_size;
19
  float q_fragment[HeadDim / 32];
20
  float o_fragment[HeadDim / 32] = {};
21
  load_query_fragment(q, batch, query_head, lane, q_fragment);
22

23
  float running_max = -CUDART_INF_F;
24
  float running_sum = 0.0F;
25
  for (int token = 0; token < lengths[batch]; ++token) {
26
    const int logical_page = token / PageSize;
27
    const int offset = token % PageSize;
28
    const int page = page_table[batch * max_pages + logical_page];
29
    const half* key = cache_ptr(cache, page, 0, offset, kv_head);
30
    const half* value = cache_ptr(cache, page, 1, offset, kv_head);
31

32
    float score = dot_fragment<HeadDim>(q_fragment, key, lane);
33
    score = warp_sum(score) * rsqrtf(static_cast<float>(HeadDim));
34
    const float next_max = fmaxf(running_max, score);
35
    const float old_scale = __expf(running_max - next_max);
36
    const float weight = __expf(score - next_max);
37
    #pragma unroll
38
    for (int i = 0; i < HeadDim / 32; ++i) {
39
      o_fragment[i] = o_fragment[i] * old_scale
40
                    + weight * __half2float(value[lane + 32 * i]);
41
    }
42
    running_sum = running_sum * old_scale + weight;
43
    running_max = next_max;
44
  }
45
  store_output(out, batch, query_head, lane, o_fragment, running_sum);
46
}