Introduction

During recent interviews, I have been asked a question about my experience with optimizing
complex kernels or experience on architecture higher than sm89. Unfortunately, I don’t have much experience on this topic, but a few days ago, I have earned a H100x8 server for some reason, so I deside to spend a weekend to optimize my decode paged attention kernel in NanoPD. The post is the first part of the optimization process, which contains my reading notes on the vLLM and flash infer kernels.

VLLM’s paged attention kernel

First let us admire the vLLM’s paged attention kernel, which is implemented in this file. In the following part, I will line by line analyze the code and try to understand the optimization techniques used in this kernel.

Block sum kernel

Ahead of directly analysing the paged attention kernel, we need to understand the block sum kernel:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

template < int NUM_WARPS> 
inline __device__ float block_sum(float* red_smem, float sum) { 
 // Decompose the thread index into warp / lane.
 int warp = threadIdx.x / WARP_SIZE; 
 int lane = threadIdx.x % WARP_SIZE; 
 // Compute the sum per warp.
#pragma unroll
 for (int mask = WARP_SIZE / 2; mask > = 1; mask /= 2) { 
 sum += VLLM_SHFL_XOR_SYNC(sum, mask); 
 } 
 // Warp leaders store the data to shared memory.
 if (lane == 0) { 
 red_smem[warp] = sum; 
 } 
 // Make sure the data is in shared memory.
 __syncthreads(); 
 // The warps compute the final sums.
 if (lane < NUM_WARPS) { 
 sum = red_smem[lane]; 
 } 
 // Parallel reduction inside the warp.
#pragma unroll
 for (int mask = NUM_WARPS / 2; mask > = 1; mask /= 2) { 
 sum += VLLM_SHFL_XOR_SYNC(sum, mask); 
 } 
 // Broadcast to other threads.
 return VLLM_SHFL_SYNC(sum, 0); 
} 

First reduce in a warp, then store the result in shm, then reduce in the warp again, which is a common reduction pattern. No bank conflict as Warp Size is 32 or less, and the shared memory is float array, so each element is 4 bytes, which means each warp will write to a different bank. The reduction in the warp is done by shuffle instructions, which can be very efficient. But warp divergence can happen in the reduction, but it is not a big problem because the reduction is done in log2(NUM_WARPS) steps, which is small. The final result is broadcasted to all threads in the block, which can be efficient if NUM_WARPS is small.

Paged attention kernel

Now the main dish comes, first let us look at the kernel signature:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

// Grid: (num_heads, num_seqs, max_num_partitions).
template < typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE, 
 int NUM_THREADS, vllm: :Fp8KVCacheDataType KV_DTYPE, 
 bool IS_BLOCK_SPARSE, 
 int PARTITION_SIZE = 0> // Zero means no partitioning.
__device__ void paged_attention_kernel(
 float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
 float* __restrict__ max_logits, // [num_seqs, num_heads, 
 // max_num_partitions]
 scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, 
 // head_size]
 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
 const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, 
 // head_size/x, block_size, x]
 const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, 
 // head_size, block_size]
 const int num_kv_heads, // [num_heads]
 const float scale, 
 const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
 const int* __restrict__ seq_lens, // [num_seqs]
 const int max_num_blocks_per_seq, 
 const float* __restrict__ alibi_slopes, // [num_heads]
 const int q_stride, const int kv_block_stride, const int kv_head_stride, 
 const float* k_scale, const float* v_scale, const int tp_rank, 
 const int blocksparse_local_blocks, const int blocksparse_vert_stride, 
 const int blocksparse_block_size, const int blocksparse_head_sliding_step)

First consider the template parameters:

1
2
3
4
5
6
7
8

typename scalar_t; // Q/K/V data type, can be float16, bfloat16, int8, etc.
typename cache_t; // KV cache store data type, can be float16, bfloat16, used for quant
int HEAD_SIZE; // head size, can be 64, 128, 256, etc.
int BLOCK_SIZE; // the block size of paged attention, can be 128, 256, etc.
int NUM_THREADS; // the number of threads per block, can be 128, 256
vllm: :Fp8KVCacheDataType KV_DTYPE; // kv quantization data type, can be int8, int4, etc.
bool IS_BLOCK_SPARSE; // whether the attention is block sparse, if true, the block_tables will be used to determine which blocks are valid.
int PARTITION_SIZE; // the partition size of the attention.

Partition attention:
At the decode stage, the query may be fewer than the prefill stage, but we still need to tackle a sequence-long kv cache. So the idea is to cut the kv cache into multiple partitions, and each partition will be processed by one kernel. Thus we achieved better parallelism and memory access pattern. The partition size can be tuned for better performance, and it is usually set to be a multiple of the block size.(from Tri Daos’ 2023 work Flash-Decoding for long-context inference)

Then consider the grid and block configuration:
The comment said Grid:(num_heads, num_seqs, max_num_partitions):

1
2
3

blockIdx.x: head index, range from 0 to num_heads-1
blockIdx.y: sequence index, range from 0 to num_seqs-1
blockIdx.z: partition index, range from 0 to max_num_partitions-1

Each cuda block will process one sequence of one head for one partition.

Output parameters:

1
2
3

float * __restrict__ exp_sums; // [num_seqs, num_heads, max_num_partitions]
float * __restrict__ max_logits; // [num_seqs, num_heads, max_num_partitions]
scalar_t * __restrict__ out; // [num_seqs, num_heads, max_num_partitions, head_size]

They are the middle results of Flash-Decoding. When PARTITION_SIZE > 0, each partition computes local softmax, expsum, weighted ouput and then reduce them to get the final output. PARTITION_SIZE = 0 means no partition, the kernel will compute the final output directly without reduction.

Input parameters:

1
2
3
4
5
6
7
8
9
10
11
12
13

scalar_t * __restrict__ q; // [num_seqs, num_heads, head_size]
cache_t * __restrict__ k_cache; // [num_blocks, num_kv_heads, head_size/x, block_size, x]
cache_t * __restrict__ v_cache; // [num_blocks, num_kv_heads, head_size, block_size]
int num_kv_heads; 
float scale; 
int * __restrict__ block_tables; // [num_seqs, max_num_blocks_per_seq]
int * __restrict__ seq_lens; // [num_seqs]
int max_num_blocks_per_seq; 
float * __restrict__ alibi_slopes; // [num_heads]
int q_stride, kv_block_stride, kv_head_stride; 
float *k_scale, *v_scale; 
int tp_rank; 
int blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size, blocksparse_head_sliding_step; 

• q: the query tensor, each sequence has one query vector for each head, so the shape is [num_seqs, num_heads, head_size].
• k_cache and v_cache: the kv cache tensor, each sequence has multiple blocks of kv cache, each block has multiple kv heads, each kv head has multiple key/value vectors, so the shape is [num_blocks, num_kv_heads, head_size/x, block_size, x], where we split head_size into x parts to vectorize the memory access for better performance. The actual head size is head_size/x * x = head_size, and the actual block size is block_size * x.
• num_kv_heads: the number of kv heads for each head.
• scale: the scaling factor for the attention, usually set to be 1/sqrt(head_size).
• block_tables: the block table for block sparse attention, each sequence has a block table to indicate which blocks are valid, so the shape is [num_seqs, max_num_blocks_per_seq].
• seq_lens: the actual sequence length for each sequence, so the shape is [num_seqs].
• max_num_blocks_per_seq: the maximum number of blocks for each sequence, used for block sparse attention.
• alibi_slopes: the alibi slopes for each head when we don’t use RoPE, we do not discuss it here, so just ignore it.
• q_stride, kv_block_stride, kv_head_stride: the stride for accessing the q, k_cache and v_cache tensors, avoid calculating the stride in the kernel for better performance.
• k_scale and v_scale: the scaling factor for the quantized k and v, used for unquantization.
• tp_rank: the tensor parallel rank of the current process, used for partitioning the kv cache for tensor parallelism.
• blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size, blocksparse_head_sliding_step: the parameters for block sparse attention, used for calculating the valid blocks for each head.

Next comes the kernel body, first to decide the work space of now block:

1
2
3
4
5

const int seq_idx = blockIdx.y; 
const int partition_idx = blockIdx.z; 
const int max_num_partitions = gridDim.z; 
constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0; 
const int seq_len = seq_lens[seq_idx]; 

Read the sequence index and partition index from the block index, and read the sequence length from the seq_lens tensor.

Then calculate the start and end position of the current partition:

1

const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); 

Calculate how many blocks are there for the current sequence, which is the sequence length divided by the block size, rounded up.

1
2

const int num_blocks_per_partition =
 USE_PARTITIONING? PARTITION_SIZE / BLOCK_SIZE: num_seq_blocks; 

Calculate how many blocks are there for each partition.

1
2
3
4
5

const int start_block_idx =
 USE_PARTITIONING? partition_idx * num_blocks_per_partition: 0; 
const int end_block_idx =
 MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks); 
const int num_blocks = end_block_idx - start_block_idx; 

Obviously, we can infer from the code itself.

Then convert the block range to the token range:

1
2
3

const int start_token_idx = start_block_idx * BLOCK_SIZE; 
const int end_token_idx = MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len); 
const int num_tokens = end_token_idx - start_token_idx; 

Thread Group Design

1

constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1); 

A thread group is a group of threads that work together to process the QK product for one token.

One warp has 32 threads and a KV block has BLOCK_SIZE tokens, so the design is to let threads balance the workload of one block. Each token receives WARP_SIZE / BLOCK_SIZE threads to compute the QK product.

1

constexpr int NUM_THREAD_GROUPS = NUM_THREADS / THREAD_GROUP_SIZE; 

Calculate how many thread groups are there in one block, which is the total number token can be processed in parallel.

1

constexpr int NUM_TOKENS_PER_THREAD_GROUP = DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE); 

If BLOCK_SIZE > WARP_SIZE, each thread group will process multiple tokens.

1
2
3
4

constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE; 
const int thread_idx = threadIdx.x; 
const int warp_idx = thread_idx / WARP_SIZE; 
const int lane = thread_idx % WARP_SIZE; 

Standard thread index calculation.

GQA head projection

1
2
3
4

const int head_idx = blockIdx.x; 
const int num_heads = gridDim.x; 
const int num_queries_per_kv = num_heads / num_kv_heads; 
const int kv_head_idx = head_idx / num_queries_per_kv; 

GQA means multi Q head with shared K/V, num_queries_per_kv is the number of query heads that share the same kv head, so we can calculate the kv head index from the query head index.

Vector Type Definition

1

constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1); 

The goal is to vectorize the memory access in a thread group to read 16 bytes of data which targeting LDG.128 instruction.

For an example: THREAD_GROUP_SIZE = 4, scalar_t = float16, then VEC_SIZE = 16 / (4 * 2) = 2, which means each thread will read 2 float16 elements once, which is 4 bytes, and the whole thread group will read 16 bytes once.

Then we can define the vector type for the q/k/v:

1
2
3

using K_vec = typename Vec< scalar_t, VEC_SIZE> : :Type; 
using Q_vec = typename Vec< scalar_t, VEC_SIZE> : :Type; 
using Quant_vec = typename Vec< cache_t, VEC_SIZE> : :Type; 

Then we can calculate how many elements in one thread:

1
2

constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE; 
constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE; 

HEAD_SIZE elements is assigned to one thread group, each thread process NUM_ELEMS_PER_THREAD elements. Then make groups according to the VEC_SIZE for vectorized memory access, each thread will process NUM_VECS_PER_THREAD vectors.

Coordinate in the Thread Group

1
2

const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE; 
const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE; 

Calculate the thread group index and the offset of the thread in the thread group, which will be used for memory access and reduction.

Load the Query to registers.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

 // Load the query to registers.
 // Each thread in a thread group has a different part of the query.
 // For example, if the thread group size is 4, then the first thread in
 // the group has 0, 4, 8, ... th vectors of the query, and the second thread
 // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
 // q is split from a qkv tensor, it may not be contiguous.
 const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE; 
 __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD]; 
#pragma unroll
 for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD; 
 i += NUM_THREAD_GROUPS) { 
 const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE; 
 q_vecs[thread_group_offset][i] =
 *reinterpret_cast< const Q_vec*> (q_ptr + vec_idx * VEC_SIZE); 
 } 
 __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a
 // memory wall right before we use q_vecs

Q‘s shape is [num_seqs, num_heads, head_size] and the grid shape is [num_heads, num_seqs, max_num_partitions], so the query for the current block can be calculated by q +seq_idx * q_stride + head_idx * HEAD_SIZE. The shared memory q_vecs layout is suitable for the thread group design.

Then each thread in the thread group(assume it’s size is 4) will traverse a row of the shm, thread 0, 4, 8, 12 will traverse the first row, then their thread group idx is 0, 1, … NUM_THREAD_GROUPS - 1, we assign them to a row of the shm and repeat it. For example, if NUM_THREAD_GROUPS = 4, then thread 0, 4, 8, 12 will traverse(0, 1, 2, 3), (4, 5, 6, 7)…
The Q access model is computed by vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE, meaning that the first row of shm stores (0, 4, 8, 12), the second row stores (1, 5, 9, 13), and so on. This access pattern can ensure coalesced memory access for the query.

Why we store the Query in shared memory instead of registers?
The query will be reused for multiple thread groups. After storing it into shm, each thread can access the query of other thread groups for free.

Shared Memory Plan

1
2
3

extern __shared__ char shared_mem[]; 
float* logits = reinterpret_cast< float*> (shared_mem); 
__shared__ float red_smem[2 * NUM_WARPS]; 

The shared_mem is used to store the attention score of all the tokens and the red_smem is used for the reduction.

1

constexpr int x = 16 / sizeof(cache_t); 

K cache’s layout is [num_blocks, num_kv_heads, head_size/x, block_size, x], so we need to calculate the x for vectorized memory access.

1

float qk_max = -FLT_MAX; 

Each thread keep a local max used for online softmax.

1

const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; 

block_tables‘s layout is [num_seqs, max_num_blocks_per_seq], we get the block table information.

Here we would not consider the block sparse case, so just ignore the block sparse related code for now.

Main Loop

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; 
 block_idx += NUM_WARPS) { 
 // NOTE(woosuk): The block number is stored in int32. However, we cast it to
 // int64 because int32 can lead to overflow when this variable is multiplied
 // by large numbers (e.g., kv_block_stride).
 // For blocksparse attention: skip computation on blocks that are not
 // attended
 if constexpr (IS_BLOCK_SPARSE) { 
 const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; 
 const bool is_remote =
 ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0); 
 const bool is_local =
 (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks); 
 if (!is_remote & & !is_local) { 
 for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { 
 const int physical_block_offset =
 (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; 
 const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; 

 if (thread_group_offset == 0) { 
 // NOTE(linxihui): assign very large number to skipped tokens to
 // avoid contribution to the sumexp softmax normalizer. This will
 // not be used at computing sum(softmax*v) as the blocks will be
 // skipped.
 logits[token_idx - start_token_idx] = -FLT_MAX; 
 } 
 } 
 continue; 
 } 
 } 
 const int64_t physical_block_number =
 static_cast< int64_t> (block_table[block_idx]); 

 // Load a key to registers.
 // Each thread in a thread group has a different part of the key.
 // For example, if the thread group size is 4, then the first thread in
 // the group has 0, 4, 8, ... th vectors of the key, and the second thread
 // has 1, 5, 9, ... th vectors of the key, and so on.
 for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { 
 const int physical_block_offset =
 (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; 
 const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; 
 K_vec k_vecs[NUM_VECS_PER_THREAD]; 

#pragma unroll
 for (int j = 0; j < NUM_VECS_PER_THREAD; j++) { 
 const cache_t* k_ptr =
 k_cache + physical_block_number * kv_block_stride +
 kv_head_idx * kv_head_stride + physical_block_offset * x; 
 const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE; 
 const int offset1 = (vec_idx * VEC_SIZE) / x; 
 const int offset2 = (vec_idx * VEC_SIZE) % x; 

 if constexpr (KV_DTYPE == Fp8KVCacheDataType: :kAuto) { 
 k_vecs[j] = *reinterpret_cast< const K_vec*> (
 k_ptr + offset1 * BLOCK_SIZE * x + offset2); 
 } else { 
 // Vector conversion from Quant_vec to K_vec.
 Quant_vec k_vec_quant = *reinterpret_cast< const Quant_vec*> (
 k_ptr + offset1 * BLOCK_SIZE * x + offset2); 
 k_vecs[j] = fp8: :scaled_convert< K_vec, Quant_vec, KV_DTYPE> (
 k_vec_quant, *k_scale); 
 } 
 } 

 // Compute dot product.
 // This includes a reduction across the threads in the same thread group.
 float qk = scale * Qk_dot< scalar_t, THREAD_GROUP_SIZE> : :dot(
 q_vecs[thread_group_offset], k_vecs); 
 // Add the ALiBi bias if slopes are given.
 qk += (alibi_slope!= 0)? alibi_slope * (token_idx - seq_len + 1): 0; 

 if (thread_group_offset == 0) { 
 // Store the partial reductions to shared memory.
 // NOTE(woosuk): It is required to zero out the masked logits.
 const bool mask = token_idx > = seq_len; 
 logits[token_idx - start_token_idx] = mask? 0.f: qk; 
 // Update the max value.
 qk_max = mask? qk_max: fmaxf(qk_max, qk); 
 } 
 } 
 } 

KV block assign loop

1
2

for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; 
 block_idx += NUM_WARPS) 

Each warp tackle a few KV blocks, if there are 4 warps and 16 blocks, then the warp 1 will tackle block 1, 5, 9, 13.

Physical block location

1
2

const int64_t physical_block_number =
 static_cast< int64_t> (block_table[block_idx]); 

The block table is used to map the logical block index to the physical block index in the KV cache, which is used for block sparse attention.

Inner loop assign token in a warp

1
2
3
4

for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { 
 const int physical_block_offset =
 (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; 
 const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; 

Usually, NUM_TOKENS_PER_THREAD_GROUP is 1, which means each thread group process one token.

Load K

1
2
3
4
5
6
7
8
9

K_vec k_vecs[NUM_VECS_PER_THREAD]; 

for (int j = 0; j < NUM_VECS_PER_THREAD; j++) { 
 const cache_t* k_ptr =
 k_cache + physical_block_number * kv_block_stride +
 kv_head_idx * kv_head_stride + physical_block_offset * x; 
 const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE; 
 const int offset1 = (vec_idx * VEC_SIZE) / x; 
 const int offset2 = (vec_idx * VEC_SIZE) % x; 

K cache’s layout is [num_blocks, num_kv_heads, head_size/x, block_size, x], so the base ptr is located to [physical_block_number, kv_head_idx, 0, physical_block_offset, 0]. This is the start position of the k vector for the current token. Then the vec_idx is used like the Q loading part to ensure thread_group_offset thread process the vec_idxth vector.

1
2

k_vecs[j] = *reinterpret_cast< const K_vec*> (
 k_ptr + offset1 * BLOCK_SIZE * x + offset2); 

Then we can load the k vector to registers.
Assume THREAD_GROUP_SIZE=4, VEC_SIZE=2, x=8, NUM_VECS_PER_THREAD=8, a thread group will process a token as below:

1
2
3
4
5
6
7
8

thread_group_offset=0, j=0: vec_idx=0, offset1=0, offset2=0 → [0,1]
thread_group_offset=1, j=0: vec_idx=1, offset1=0, offset2=2 → [2,3]
thread_group_offset=2, j=0: vec_idx=2, offset1=0, offset2=4 → [4,5]
thread_group_offset=3, j=0: vec_idx=3, offset1=0, offset2=6 → [6,7]

thread_group_offset=0, j=1: vec_idx=4, offset1=1, offset2=0 → [8,9]
thread_group_offset=1, j=1: vec_idx=5, offset1=1, offset2=2 → [10,11]
...

Dot product and store the logits

1
2
3
4

// Compute dot product.
// This includes a reduction across the threads in the same thread group.
float qk = scale * Qk_dot< scalar_t, THREAD_GROUP_SIZE> : :dot(
 q_vecs[thread_group_offset], k_vecs); 

Each thread compute the dot product and reduce in the thread group by Qk_dot, which is implemented by warp shuffle.

1
2
3
4
5

if (thread_group_offset == 0) { 
 const bool mask = token_idx > = seq_len; 
 logits[token_idx - start_token_idx] = mask? 0.f: qk; 
 qk_max = mask? qk_max: fmaxf(qk_max, qk); 
} 

Then store the logits to shared memory, and update the local max for softmax.

Reduction for max logits

1
2
3
4
5
6
7
8
9
10
11

// Perform reduction across the threads in the same warp to get the
// max qk value for each " warp" (not across the thread block yet).
// The 0-th thread of each thread group already has its max qk value.
#pragma unroll
 for (int mask = WARP_SIZE / 2; mask > = THREAD_GROUP_SIZE; mask /= 2) { 
 qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); 
 } 
 if (lane == 0) { 
 red_smem[warp_idx] = qk_max; 
 } 
 __syncthreads(); 

Reduce all the kv block processed by a warp.

1
2
3
4
5
6
7

 qk_max = lane < NUM_WARPS? red_smem[lane]: -FLT_MAX; 
#pragma unroll
 for (int mask = NUM_WARPS / 2; mask > = 1; mask /= 2) { 
 qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); 
 } 
 // Broadcast the max qk value to all threads.
 qk_max = VLLM_SHFL_SYNC(qk_max, 0); 

Then cross warps to get the max qk value.

Then compute the exp and the exp sums:

1
2
3
4
5
6
7
8

float exp_sum = 0.f; 
for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { 
 float val = __expf(logits[i] - qk_max); 
 logits[i] = val; 
 exp_sum += val; 
} 

exp_sum = block_sum< NUM_WARPS> (& red_smem[NUM_WARPS], exp_sum); 

After this, we normalize the logits and get the softmax output:

1
2
3
4
5

const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f); 
for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { 
 logits[i] *= inv_sum; 
} 
__syncthreads(); 

Load V into the memory

V cache’s layout is [num_blocks, num_kv_heads, head_size, block_size], which is different from K cache.

1
2
3
4
5
6
7
8
9
10
11
12
13

constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE); 
using V_vec = typename Vec< scalar_t, V_VEC_SIZE> : :Type; 
using L_vec = typename Vec< scalar_t, V_VEC_SIZE> : :Type; 
using V_quant_vec = typename Vec< cache_t, V_VEC_SIZE> : :Type; 
using Float_L_vec = typename FloatVec< L_vec> : :Type; 

constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE; 
constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW; 
constexpr int NUM_ROWS_PER_THREAD =
 DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER); 

// NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
float accs[NUM_ROWS_PER_THREAD]; 

There are BLOCK_SIZE elements in one row of V, split them in V_VEC_SIZE. A warp has WARP_SIZE threads, so one iteration can process NUM_ROWS_PER_ITER rows, and each thread will process NUM_ROWS_PER_THREAD rows.

Acc initialization

1
2
3
4

// Initialize the accumulators.
for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { 
 accs[i] = 0.f; 
} 

Then comes the second main part: the matrix multiplication between the softmax output and the V vectors.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

scalar_t zero_value; 
 zero(zero_value); 
 for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; 
 block_idx += NUM_WARPS) { 
 // NOTE(woosuk): The block number is stored in int32. However, we cast it to
 // int64 because int32 can lead to overflow when this variable is multiplied
 // by large numbers (e.g., kv_block_stride).
 // For blocksparse attention: skip computation on blocks that are not
 // attended
 if constexpr (IS_BLOCK_SPARSE) { 
 int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; 
 if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) & & 
!((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) { 
 continue; 
 } 
 } 
 const int64_t physical_block_number =
 static_cast< int64_t> (block_table[block_idx]); 
 const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE; 
 const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; 
 L_vec logits_vec; 
 from_float(logits_vec, *reinterpret_cast< Float_L_vec*> (logits + token_idx -
 start_token_idx)); 

 const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
 kv_head_idx * kv_head_stride; 
#pragma unroll
 for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { 
 const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; 
 if (row_idx < HEAD_SIZE) { 
 const int offset = row_idx * BLOCK_SIZE + physical_block_offset; 
 V_vec v_vec; 

 if constexpr (KV_DTYPE == Fp8KVCacheDataType: :kAuto) { 
 v_vec = *reinterpret_cast< const V_vec*> (v_ptr + offset); 
 } else { 
 V_quant_vec v_quant_vec =
 *reinterpret_cast< const V_quant_vec*> (v_ptr + offset); 
 // Vector conversion from V_quant_vec to V_vec.
 v_vec = fp8: :scaled_convert< V_vec, V_quant_vec, KV_DTYPE> (v_quant_vec, 
 *v_scale); 
 } 
 if (block_idx == num_seq_blocks - 1) { 
 // NOTE(woosuk): When v_vec contains the tokens that are out of the
 // context, we should explicitly zero out the values since they may
 // contain NaNs. See
 // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
 scalar_t* v_vec_ptr = reinterpret_cast< scalar_t*> (& v_vec); 
#pragma unroll
 for (int j = 0; j < V_VEC_SIZE; j++) { 
 v_vec_ptr[j] = token_idx + j < seq_len? v_vec_ptr[j]: zero_value; 
 } 
 } 
 accs[i] += dot(logits_vec, v_vec); 
 } 
 } 
 } 

This code is similar to the previous loop for K.

Introduction

I have learned CUDA kernel programming for a long time, but I have never learnt CUDA streams, only knowing that CUDA streams can be used to achieve concurrency. Today by reading the NVIDIA slides on CUDA streams, I have a better understanding of CUDA streams and concurrency.

Default stream

By default, all CUDA operations are issued into a single stream, called the default stream. Operations in the default stream are executed sequentially, and they are not concurrent with any other operations. The special behavior of the default stream is that it is wholely sync for host and device, which means each time we submit a operation to the default stream, the host will insert an implicit cudaDeviceSynchronize() after and before the operation. But there are several exceptions to this rule:

• Kernel launches in the default stream are asynchronous with respect to the host, but they are still serialized with respect to other operations in the default stream.
• cudaMemcpyAsync() and cudaMemsetAsync() operations in the default stream are asynchronous with respect to the host.
• cudaMemcpy() in the same device.
• cudaMemcpy() below 64KB between host and device.

Requirements for Concurrency

To achieve concurrency, we need to meet the following requirements:

• Use non-default streams for concurrent operations.
• cudaMemcpyAsync() with host from pinned memory.
• sufficient resources must be available on the device to execute concurrent operations.

Some examples

1
2
3
4
5
6
7
8

cudaMalloc(& dev1, size); 
double * host1 = (double *)malloc(& host, size); 
...
cudaMemcpy(dev1, host1, size, cudaMemcpyHostToDevice); 
kernel2<<< grid, block, 0>>> (..., dev2, ...); 
kernel3<<< grid, block, 0>>> (..., dev3, ...); 
cudaMemcpy(host4, dev4, size, cudaMemcpyDeviceToHost); 

Above code will be executed synchronously, because all operations are issued into the default stream. Observing the nsys timeline, we can see that all operations are executed sequentially.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

cudaStream_t streams[NUM_STREAMS]; 
for(int i = 0; i < NUM_STREAMS; i ++){ 
 cudaStreamCreate(& streams[i]); 
} 
for(int i = 0; i < NUM_STREAMS; i ++){ 
 int offset = i * chunk; 
 cudaMemcpyAsync(dev1 + offset, host + offset, chunk_size, cudaMemcpyHostToDevice, streams[i]); 
} 
for(int i = 0; i < NUM_STREAMS; i ++){ 
 int offset = i * chunk; 
 kernel1<<< (chunk + 255) / 256, 256, 0, streams[i]>>> (dev1 + offset, dev2 + offset, chunk); 
} 
for(int i = 0; i < NUM_STREAMS; i ++){ 
 int offset = i * chunk; 
 cudaMemcpyAsync(host + offset, dev2 + offset, chunk_size, 
 cudaMemcpyDeviceToHost, streams[i]); 
} 

Above code will be executed concurrently, because we have issued operations into different streams. Observing the nsys timeline, we can see that all operations are executed concurrently.

Another overlap example is as follows:

1
2
3
4
5

cudaMemcpy(dev1, host1, size, H2D); 
kernel2<<< grid, block>>> (dev2); // launch kernel is asynchronous with respect to the host.
some_CPU_method(); // overlap with kernel2
kernel3<<< grid, block>>> (dev3); 
cudaMemcpy(host4, dev4, size, D2H); 

In above code, kernel2 will be launched asynchronously with respect to the host, so some_CPU_method() can be executed concurrently with kernel2. However, kernel3 and cudaMemcpy() will be executed sequentially after kernel2, because they are issued into the default stream.

Explicit Synchronization

• Synchronize everything: cudaDeviceSynchronize(): blocks host until all issued CUDA operations are completed.
• Synchronize a stream: cudaStreamSynchronize(stream): blocks host until all operations in the specified stream are completed.
• Synchronize using Events: cudaEventSynchronize(event): blocks host until the specified event is completed. Events can be used to measure the time between operations in different streams.

Some Event Using Examples

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

cudaEvent_t start, stop; 
cudaEventCreate(& start); 
cudaEventCreate(& stop); 

cudaMemcpyAsync(dev1, host1, size, H2D, stream1); 
cudaEventRecord(start, stream1); // record start event after memcpy

cudaMemcpyAsync(host2, dev2, size, D2H, stream2); 
cudaStreamWaitEvent(stream2, start, 0); // make stream2 wait for the start event
kernel<<< grid, block, 0, stream2>>> (...); // kernel will execute after the start event is recorded
cudaEventRecord(stop, stream2); // record stop event after kernel launch
cudaEventSynchronize(stop); // wait for the stop event to complete
float elapsedTime; 
cudaEventElapsedTime(& elapsedTime, start, stop); // calculate elapsed time between start and stop events
printf(" Elapsed time: %f ms\n" , elapsedTime); 

Implicit Synchronization

Some operations will cause implicit synchronization, without knowing it, we may introduce unexpected synchronization points in our code, which can lead to performance degradation. Some examples of implicit synchronization are as follows:

• cudaMallocHost()/cudaFreeHost(): These functions will block the host until all previously issued CUDA operations are completed, because they need to ensure that the pinned memory is not being used by any ongoing CUDA operations.
• cudaMalloc(): This function will block the host until all previously issued CUDA operations are completed, because it needs to ensure that there are sufficient resources available on the device to allocate the requested memory.
• cudaMemcpy(): This function needs to ensure that the data transfer is not being interfered by any ongoing CUDA operations.
• cudaDeviceSetCacheConfig(): This function needs to ensure that the cache configuration is not being changed while any ongoing CUDA operations are using the cache, so it will block the host until all previously issued CUDA operations are completed.

The right way to avoid implicit synchronization is to assign the memory allocation and deallocation in the beginning and the end of the program, and use cudaMemcpyAsync() instead of cudaMemcpy() for data transfer between host and device.

Stream Scheduling

Take Fermi architecture as an example, it has 3 queues: 1 compute engine queue, 2 copy engine queues (one for H2D and one for D2H).

The shedule rule is as follows:
CUDA operations are pushed into the target queue based on the type of operation in the launch order. One operation is issued only when the three conditions are met:

• In the same stream, all previously issued operations have been completed.
• Ahead of the operation in the same queue, there is no other operation that is still executing.
• The resources required for the operation are available on the device.

One blocked operation can block the entire queue even there are other operations in the queue belonging to different streams. So the launch order of operations can affect the performance of the program.

An example of stream scheduling is as follows:

Concurrent Kernel Scheduling

Normally, a signal is inserted into the queues, after the operation is issued, to indicate the completion of the operation. But for the compute engine queues, when compute kernels are issued sequentially, the signal is not inserted until the kernel is completed. So if there are multiple kernels issued into the compute engine queue, they will be executed sequentially, even if they belong to different streams.

In some situations this delay of signals can block other queues.

Conclusion

Maybe the slides I read is a bit old, but it still gives me a good understanding of CUDA streams and concurrency. I will try to use CUDA streams in my future projects to improve the performance of my code.

Introduction

In the previous post, we have introduced the estimate of value function: MC and TD. In this post, we will introduce two important algorithms for estimating the action value function: SARSA and Q-learning.

Looking back at our previous post, now we have known ‘What is the best state’: estimating the state value function $V^{\pi }(S_t)$, but we still don’t know ‘What is the best action’: $\pi (s)=\arg {\max }_{a\in A}P(s’|s,a)V^{\pi }(s’)$. Here we don’t know the transition probability $P(s’|s,a)$, so we can’t directly compute the optimal policy. So, we need to estimate the action value function $Q^{\pi }(s,a)$.

SARSA

For any (state, action, reward, next state, next action) executated by the policy $\pi$, we can update the action value function as follows:$$Q(S_t,A_t)\leftarrow Q(S_t,A_t)+\alpha \left[R_{t+1}+\gamma Q(S_{t+1},A_{t+1})-Q(S_t,A_t)\right]$$

Introduction

In the previous posts, we have introduced the MDP and its solution. But in practice, we often do not have the simulation of the environment, which means we cannot directly apply our knowledge of MDP to solve the problem. There is indeed a method to simulate the environment, which is called model-based Reinforcement learning. However, in this post, we will focus on the model-free Reinforcement learning, which does not require the simulation of the environment and the construction of MDPs.

Estimating Value Functions

In mode-based RL, value functions can be computed by DP methods as follows:$$\begin{array}{rl}{V}^{\pi }(s)& =E_{\pi }\left[R(s_0,a_0)+\gamma R(s_1,a_1)+{\gamma }^{2}R(s_2,a_2)+\cdots |s_0=s\right]\\ & =R(s,a)+\gamma \sum _{s’}P(s’|s,a)V^{\pi }(s’)\end{array}$$

However, in model-free RL, we cannot directly access the $P(s’|s,a)$ and $R(s,a)$, but we have some ways to estimate the value function from episodes of experience.

Why we estimate the value function?
Because we can use the value function to derive the optimal policy, which is our ultimate goal. Besides value function can help us to reuse historical experience to make better decisions in the future, which is the essence of Reinforcement learning.

Here is a graph to introduce some methods to estimate the value function:

Monte Carlo methods

Target: Learn $V^{\pi }$ from episodes of experience.

Review: accumulate reward function:

$$G_t=R_t+\gamma R_{t+1}+{\gamma }^2{R}_{t+2}+\cdots =\sum _{k=0}^{\infty }{\gamma }^k{R}_{t+k}$$

Review: value function is the expected return:$$V^{\pi }(s)=E_{\pi }\left[G_t|s_t=s\right]\simeq \frac{1}{N(s)}\sum _{i=1}^{N(s)}{G}_t^i$$
The Monte Carlo method use empirical mean cumulative reward instead of expected return to estimate the value function.

First-visit Monte Carlo method

The first-visit Monte Carlo method estimates the value function by averaging the returns following the first time a state is visited in an episode. The algorithm is as follows:

1. Initialization:
   • For any $s\in S$,$V(s)\in \mathbb{R}$, $N(s)=0$.
   • For any $s\in S$,$\text{returns}(s)=\varnothing$.
2. Loop for each episode:
   • Generate an episode following policy $\pi$: $S_0,A_0,R_1,S_1,A_1,R_2,\cdots ,S_T$.
   • For t = T-1, T-2, …, 0:
     • $G\leftarrow \gamma G+R_{t+1}$
     • If $s$ is the first time in the episode:
       • Append $G$ to $\text{returns}(s)$.
       • $V(s)\leftarrow \text{average}(\text{returns}(s))$.

The reason why we call it ‘first-visit’ is that we only update the value function for the first time we visit a state in an episode. Thus we can avoid the bias caused by multiple visits to the same state in an episode. However, this method may have high variance because it only uses one return for each state in an episode.

Incremental Monte Carlo method

The first-visit Monte Carlo method takes a lot of memory to store the returns for each state, which is not efficient. The incremental Monte Carlo method uses an incremental update rule to estimate the value function without storing all the returns. The algorithm is as follows:

1. Initialization:
   • For any $s\in S$,$V(s)\in \mathbb{R}$, $N(s)=0$, $G=0$.
2. Loop for each episode:
   • Generate an episode following policy $\pi$: $S_0,A_0,R_1,S_1,A_1,R_2,\cdots ,S_T$.
   • For t = T-1, T-2, …, 0:
     • $G\leftarrow \gamma G+R_{t+1}$
     • If $s$ is the first time in the episode:
       • $N(s)\leftarrow N(s)+1$
       • $V(s)\leftarrow V(s)+\frac{1}{N(s)}(G-V(s))$.

Interesting, online softmax also takes the same update rule as the incremental Monte Carlo method. Great job.

Besides, incremental MC provides more design space for us to tackle some problems in practice. For example, we can use a constant step size $\alpha$ instead of $\frac{1}{N(s)}$ to update the value function, which is called constant step size MC method. It is useful when the environment is non-stationary, which means the reward function and transition probability may change over time. In this case, we want to give more weight to recent returns than old returns, which can be achieved by using a constant step size $\alpha$:$$V(s)\leftarrow V(s)+\alpha (G-V(s))$$

Some properties of Monte Carlo methods

1. Monte Carlo methods are model-free, which means they do not require the knowledge of the environment’s dynamics (transition probabilities and reward function).
2. Monte Carlo methods take the simpliest approach to estimate the value function, which is to average the returns following the policy. However, this method may have high variance because it only uses one return for each state in an episode.
3. One key to note is that Monte Carlo methods can only be applied to finite MDPs, which means the state space and action space must be finite.

Importance sampling

Let’s try to estimate a custom distribution $p(x)$ ‘s expectation.$$\begin{array}{rl}{E}_{x\sim p}[f(x)]& =\int f(x)p(x)dx\\ & =\int f(x)\frac{p(x)}{q(x)}q(x)dx\\ & =E_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right]\end{array}$$
Then we reassign the importance sampling weight $w(x)=\frac{p(x)}{q(x)}$, we can rewrite the above equation as:$$E_{x\sim p}[f(x)]=E_{x\sim q}\left[f(x)w(x)\right]$$

Off-policy Monte Carlo methods via Importance Sampling

We can use the cumulative reward function of policy $\mu$ to justify policy $\pi$, and then weight the cumulative reward function by the importance ratio between $\pi$ and $\mu$ to estimate the value function of policy $\pi$. The algorithm is as follows:

Every episode would be mutified by the importance sampling ratio:$$G_t^{\pi /\mu }=\frac{\pi (A_t|S_t)}{\mu (A_t|S_t)}\frac{\pi (A_{t+1}|S_{t+1})}{\mu (A_{t+1}|S_{t+1})}\cdots \frac{\pi (A_{T-1}|S_{T-1})}{\mu (A_{T-1}|S_{T-1})}{G}_t$$

So we then update the value function by:$$V(s)\leftarrow V(s)+\frac{1}{N(s)}(G_t^{\pi /\mu }-V(s))$$

Sample by importance sampling will significantly increase the variance of the return, which is because the importance sampling ratio can be very large when $\pi$ and $\mu$ are very different.

Temporal-Difference Learning

Temporal-Difference (TD) is a method combining the MC method and DP method, which name comes from the fact that it uses the diff of estimated value function at two consecutive time steps to update the value function. There are two key ideas in TD learning: TD error and TD target.

For state value funtion $V$, after a transition from state $s$ to state $s’$ with reward $r$, the TD error is defined as:$$\delta =r+\gamma V(s’)-V(s)$$
The TD target is defined as:$$\hat{V}=r+\gamma V(s’)$$

As for the TD in Bellman expectation equation, the TD error is used in estimating the expect part.

Some details of TD learning

The simpliest TD learning algorithm is called TD(0), which updates the value function by the TD error at each time step. The key equation of TD(0) is as follows:$$V(s)\leftarrow V(s)+\alpha \delta =V(s)+\alpha (r+\gamma V(s’)-V(s))$$

Why we update like this?
The Bellman expectation equation is rewritten as:

$$E_{\pi }\left[R_{t+1}+\gamma V^{\pi }(S_{t+1})-V^{\pi }(S_t)|S_t=s\right]=0$$
That’s all. We want to make the TD error as small as possible, which means we want to make the estimated value function as close as possible to the true value function. Thus, we can use the TD error to update the value function.

The TD method introduce the bootstrapping idea, which means we use the estimated value function to update the value function. This is different from the MC method, which uses the actual return to update the value function. The bootstrapping idea can significantly reduce the variance of the return, but it may introduce bias because we are using an estimated value function to update the value function.

Contrast between TD and MC methods

They have the same goal: Learn the value function from episodes of experience. However, they have different approaches to achieve this goal. The MC method uses the actual return to update the value function, which can have high variance but no bias. The TD method uses the estimated value function to update the value function, which can have low variance but may introduce bias.

The object of TD is $R_t+\gamma V(s_{t+1})$, which is called TD target, while the object of MC is $G_t$, which is the actual return. The TD method’s error is called TD error, which is defined as $\delta =r+\gamma V(s’)-V(s)$, while the MC method’s error is defined as $G_t-V(s)$.

The strengths and limitations of TD learning and MC learning

TD method can learn until the end of an episode:

• After each step in an episode, TD method can update the value function use the former value function, which means it can learn until the end of an episode. However, MC method can only update the value function after the end of an episode, which means it cannot learn until the end of an episode.
• TD method can learn from incomplete episodes, which means it can learn from episodes that are not terminated. However, MC method can only learn from complete episodes, which means it cannot learn from episodes that are not terminated.

Tradeoff between bias and variance

Note: The real TD target uses the true $V^{\pi }$, which is unknown in practice. The actual TD target uses the current estimate $V$, introducing bias. Despite the bias, TD typically has lower variance than MC because it bootstraps from a single step rather than a full trajectory.

Multi-step TD learning

The TD(0) method only uses the immediate reward and the estimated value of the next state to update the value function, which may not be sufficient to capture the long-term dependencies in the environment. The multi-step TD learning method uses the rewards and estimated values of multiple future states to update the value function, which can better capture the long-term dependencies in the environment. We will introduce it by leading into n-step cumulate reward function and n-step TD target.

N-step cumulate reward function

Consider the following n-step cumulate reward function:$$G_t^{(n)}=R_t+\gamma R_{t+1}+\cdots +{\gamma }^{n-1}{R}_{t+n-1}+{\gamma }^{n}V(S_{t+n})$$
It seems make sense to use the n-step cumulate reward function to update the value function, which is called n-step TD learning. The key equation of n-step TD learning is as follows:$$V(s)\leftarrow V(s)+\alpha (G_t^{(n)}-V(s))$$

N-step mean cumulate reward function

Can we take up the information of different n-step cumulate reward function to update the value function? The answer is yes.
We can use a weighted average of different n-step cumulate reward functions to update the value function, which is called n-step mean TD learning. The key weight figure of weighted average is as follows:

So the n-step mean cumulate reward function is defined as:$$G_t^{\lambda }=(1-\lambda )\sum _{n=1}^{\infty }{\lambda }^{n-1}{G}_t^{(n)}$$
Then we can update the value function by:$$V(s)\leftarrow V(s)+\alpha (G_t^{\lambda }-V(s))$$

This is called TD($\lambda$) method, which is a generalization of TD(0) and MC methods. When $\lambda =0$, TD($\lambda$) reduces to TD(0) method, and when $\lambda =1$, TD($\lambda$) reduces to MC method. Thus, by adjusting the value of $\lambda$, we can control the bias-variance tradeoff in the estimation of the value function.

Conclusion about TD(λ) method

• Unless the $lambda$ is 0, TD($\lambda$) mothod is unbiased, because it’s a weighted average of unbiased n-step TD targets.

• The variance of TD($\lambda$) method is lower than that of MC method, because it uses bootstrapping to update the value function, which can reduce the variance of the return. However, the variance of TD($\lambda$) method is higher than that of TD(0) method, because it uses more rewards and estimated values to update the value function, which can increase the variance of the return.

  $$\text{Var}(aX+bY)=a^2\text{Var}(X)+b^2\text{Var}(Y)$$

• Empirically $\lambda$ is not quite commom because fast credit assignment for a given action is preferred. So MC or TD(0) is more commonly used in practice. However, TD($\lambda$) method can be useful when we want to balance the bias-variance tradeoff in the estimation of the value function, which can be achieved by adjusting the value of $\lambda$.

So TD($\lambda$) use $\lambda$ as variable while n-step TD use $n$ as variable. TD($\lambda$) is a generalization of n-step TD, which can be seen as a weighted average of infin-step TD targets. By adjusting the value of $\lambda$, we can control the bias-variance tradeoff in the estimation of the value function, which can be useful in practice when we want to balance the bias and variance in the estimation of the value function.

Conclusion

In this post, we have introduced the model-free Reinforcement learning, which does not require the simulation of the environment and the construction of MDPs. We have introduced two methods to estimate the value function from episodes of experience: Monte Carlo methods and Temporal-Difference learning. We have also introduced the n-step TD learning method, which uses the rewards and estimated values of multiple future states to update the value function, which can better capture the long-term dependencies in the environment. Finally, we have discussed the bias-variance tradeoff in the estimation of the value function, which can be controlled by adjusting the value of $\lambda$ in TD($\lambda$) method.

Introduction

In the previous post, we have introduced the Bellman equation and the linear programming formulation for MDP. In this post, we will discuss the model-based Reinforcement learning, which is a method to solve the MDP when we do not have the model of the environment.

The Settings of RL

Typically, RL is framed as MDP, exploring the enviroment and learning the optimal policy.
Generally, we can only observe the episodes and usually, we do not have the model of the environment.

So we need to introduce the model into our RL project. Model-based RL actually is a method to solve the MDP.

Dynamic Programming Based RL

Dynamic Programming for finite MDP

Our objective function is simple, just the expected return, which is defined as:$$\underset{\pi }{\max }{E}_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_{t+1}|s_0=s\right]$$

In the first episode of the series, we introduced Backward induction, which we can start from the last step and then recursively solve the problem. However, this method is not efficient for large state space, and it requires the state transition.

But meanwhile, we can also use the Bellman equation of value function to tackle the problem:$$V^{\pi }(s)=\sum _{a\in A}\pi (a|s)\left[\underset{\text{immediate reward}}{\underbrace{R(s,a)}}+\underset{\text{discount}}{\underbrace{\gamma }}\sum _{s’\in S}\underset{\text{transition}}{\underbrace{P(s’|s,a)}}\underset{\text{future value}}{\underbrace{V^{\pi }(s’)}}\right]$$

Optimal Value Function

For a state $s$, we can define the optimal value function as the maximum value function over all policies:$$V^{\ast }(s)=\underset{\pi }{\max }{V}^{\pi }(s)$$
So the optimal value function is as follows:$$V^{\ast }(s)=\underset{a\in A}{\max }R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V^{\ast }(s’)$$

So the best policy can be derived from the optimal value function as:$${\pi }^{\ast }(a|s)=\arg \underset{a\in A}{\max }R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V^{\ast }(s’)$$

for any state $s$ and policy $\pi$, there is:$$V^{\ast }(s)=V^{{\pi }^{\ast }}(s)\ge V^{\pi }(s)$$

Obviously, the value function relates to the policy, so we can iterate the optimal value function and the optimal policy until convergence. They are called Value Iteration and Policy Iteration respectively.

Value Iteration

For an MDP which is finite in both state and action space, we can use the value iteration to solve the problem. The value iteration is as follows:

1. Initialize $V(s)=0$ for all $s\in S$
2. For each state $s\in S$, update the value function as:
   $V(s)\leftarrow {\max }_{a\in A}R(s,a)+\gamma {\sum }_{s’\in S}P(s’|s,a)V(s’)$
3. Repeat step 2 until convergence.

NOTE: There isn’t any specific order to update the value function, we can update the value function in any order. But the convergence rate may be different.

Sync & Async Value Iteration

Sync value iteration need to store two copies of value funtion:

1. For any state $s$, we update the value function as:
   $V_{new}(s)\leftarrow {\max }_{a\in A}R(s,a)+\gamma {\sum }_{s’\in S}P(s’|s,a)V_{old}(s’)$
2. After updating all states, we copy the new value function to the old value function:
   $V_{old}\leftarrow V_{new}$

Async value iteration only need to store one copy of value function:
$V(s)\leftarrow {\max }_{a\in A}R(s,a)+\gamma {\sum }_{s’\in S}P(s’|s,a)V(s’)$

Policy Iteration

The assumption of MDP is the same as the value iteration, which is finite in both state and action space. The policy iteration is as follows:

1. Randomly initialize a policy $\pi$ and a value function $V(s)=0$ for all $s\in S$
2. Repeat the following steps until convergence:
3. Policy Evaluation: For each state $s\in S$, update the value function as:
   $V(s)\leftarrow {\sum }_{a\in A}\pi (a|s)\left[R(s,a)+\gamma {\sum }_{s’\in S}P(s’|s,a)V(s’)\right]$
4. Policy Improvement: For each state $s\in S$, update the policy as:
   $\pi (a|s)\leftarrow \arg {\max }_{a\in A}R(s,a)+\gamma {\sum }_{s’\in S}P(s’|s,a)V(s’)$

Obviously, the Policy Iteration will be more expensive than the Value Iteration, since it needs to evaluate the policy in each iteration. However, the Policy Iteration can converge faster than the Value Iteration, since it can update the policy in each iteration.

Let’s contrast the two methods:

NOTE:

1. Value iteration is a greedy method, we always use the best.
2. Update the value function by Bellman equation in Policy Iteration is expensive.
3. For smaller space MDP, Policy Iteration is faster than Value Iteration, but for larger space MDP, Value Iteration is faster than Policy Iteration.
4. If there isn’t any state transition circle, the value iteration is better.

Bellman operators

In fact, we have introduced the Bellman operators in the previous post, but we haven’t discussed it in detail.

Why Policy Iteration and Value Iteration can converge to the optimal value function? The key is that Bellman operators are contraction mappings.

Bellman operator is the collection of below functions:

1. Bellman expectation operator, usually denoted as ${\mathcal{T}}^{\pi }$, which is defined as:

   $${\mathcal{T}}^{\pi }V(s)=\sum _{a\in A}\pi (a|s)\left[R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V(s’)\right]$$

2. Bellman optimality operator, usually denoted as ${\mathcal{T}}^{\ast }$ or $\mathcal{T}$, which is defined as:

   $$\mathcal{T}V(s)=\underset{a\in A}{\max }R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V(s’)$$

They can be used on the state value function and action value function:

• expectation operator is used in the policy iteration, used for computing
  the value function of a given policy, while is the inner loop of the policy iteration.
• optimality operator is used in the value iteration, used for computing the optimal value function, while is the main loop of the value iteration.

Both the Bellman expectation operator and the Bellman optimality operator can be defined on the action value function and the state value function:

1. Bellman expectation operator on V-function:

$$\begin{array}{rl}{V}^{\pi }(s)& =E_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_t\mid s_0=s\right]\\ & =E_{\pi }\left[R(s_0,a_0)+\gamma \sum _{s’\in S}P(s’|s,a)\pi (a|s)V^{\pi }(s’)\right]\\ & =\sum _{a\in A}\pi (a|s)\left[R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V^{\pi }(s’)\right]\\ & =({\mathcal{T}}^{\pi }{V}^{\pi })(s)\end{array}$$

2. Bellman optimality operator on V-function:

$$\begin{array}{rl}{V}^{\ast }(s)& =\underset{\pi }{\max }{E}_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_t\mid s_0=s\right]\\ & =\underset{a\in A}{\max }R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)V^{\ast }(s’)\\ & =(\mathcal{T}{V}^{\ast })(s)\end{array}$$

3. Bellman expectation operator on Q-function:

$$\begin{array}{rl}{Q}^{\pi }(s,a)& =E_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_t\mid s_0=s,a_0=a\right]\\ & =E_{\pi }\left[R(s_0,a_0)+\gamma \sum _{s’\in S}P(s’|s,a)\pi (a|s)Q^{\pi }(s’,a’)\right]\\ & =R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)\sum _{a’\in A}\pi (a’|s’)Q^{\pi }(s’,a’)\\ & =({\mathcal{T}}^{\pi }{Q}^{\pi })(s,a)\end{array}$$

4. Bellman optimality operator on Q-function:

$$\begin{array}{rl}{Q}^{\ast }(s,a)& =\underset{\pi }{\max }{E}_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_t\mid s_0=s,a_0=a\right]\\ & =R(s,a)+\gamma \sum _{s’\in S}P(s’|s,a)\underset{a’\in A}{\max }{Q}^{\ast }(s’,a’)\\ & =(\mathcal{T}{Q}^{\ast })(s,a)\end{array}$$

Due to the contraction property of the Bellman operators, we can guarantee the convergence of the value iteration and policy iteration to the optimal value function.

Model RL

In the above sections, our objective environment is a known MDP, all our methods are based on the assumption that we have the model of the environment, which is the transition probability and the reward function. However, in many real-world scenarios, we do not have the model of the environment, so we need to learn the model from the data.

There are two basic thoughts to learn the model of the environment:

1. learn the state transition probability$P(s’|s,a)$:

   $$P(s’|s,a)=\frac{N(s,a,s’)}{N(s,a)}$$

   where $N(s,a,s’)$ is the number of times we have observed the transition from state $s$ to state $s’$ when taking action $a$, and $N(s,a)$ is the number of times we have observed taking action $a$ in state $s$.

2. learn the reward function $R(s,a)$:

   $$R(s,a)=\text{average}\left(R_t|s_t=s,a_t=a\right)$$

   where $N(s,a)$ is the number of times we have observed taking action $a$ in state $s$, and $R(s,a)$ is the average reward we have observed when taking action $a$ in state $s$.

The simple simulate algorithm is as follows:

1. randomly initialize a policy $\pi$
2. repeat the following steps until convergence:
3. collect data by executing the policy $\pi$ in the environment, and store the transition data in a replay buffer.
4. learn the model of the environment from the replay buffer, which includes learning the state transition probability and the reward function.
5. solve the MDP with the learned model to get the optimal policy ${\pi }^{\ast }$.

Other method to solve this is not learning the MDP, instead we learn the value function directly from the data, which is called model-free RL, we will discuss it in the next post.

Conclusion

In this post, we have introduced the model-based Reinforcement learning, which is a method to solve the MDP when we do not have the model of the environment. We have discussed the value iteration and policy iteration, which are two basic methods to solve the MDP. We have also introduced the Bellman operators, which are the key to guarantee the convergence of the value iteration and policy iteration. Finally, we have introduced the simple simulate algorithm, which is a method to learn the model of the environment and solve the MDP with the learned model. In the next post, we will discuss the model-free Reinforcement learning, which is a method to learn the value function directly from the data without learning the model of the environment.

Introduction

Given the former post where we have introduced the MDP and some basic properties, we are now ready to discuss the MDP-based Reinforcement learning. But first, we need to introduce the solution of MDP.

Bellman Equation

If we have learned the previous post, we will know that there are two types of value function, state value function and action value function. Their mathematical definitions are as follows:$$V^{\pi }(s)=E_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_{t+1}|s_0=s\right]$$

Instinctively, the state value function is the expected return when we start from state $s$ and follow policy $\pi$. Similarly, the action value function is defined as:

$$Q^{\pi }(s,a)=E_{\pi }\left[\sum _{t=0}^{\infty }{\gamma }^t{R}_{t+1}|s_0=s,a_0=a\right]$$

Here, it represents the expected return when we start from state $s$, take action $a$, and then follow policy $\pi$ thereafter.

On the other hand, we have a accumulate reward function, which is defined as:$$G_t=R_t+\gamma R_{t+1}+{\gamma }^2{R}_{t+2}+\cdots =\sum _{k=0}^{\infty }{\gamma }^k{R}_{t+k}$$
It can be recursively defined as:$$G_t=R_t+\gamma G_{t+1}$$

So, we can rewrite the state value function as:$$V^{\pi }(s)=\sum _a\pi (a|s)\left[R(s,a)+\gamma \sum _{s’}P(s’|s,a)V^{\pi }(s’)\right]$$
Above is the Bellman expectation equation for the state value function. Now, our question is to choose the best policy $\pi$ to maximize the value function. We can define the optimal state value function as:$$V^{\ast }(s)=\underset{\pi }{\max }{V}^{\pi }(s)$$

Due to the Principle of Optimality: Each stage of an optimal policy must be optimal for the remaining stages, we can derive the Bellman optimality equation for the state value function as:$$V^{\ast }(s)=\underset{a}{\max }\left[R(s,a)+\gamma \sum _{s’}P(s’|s,a)V^{\ast }(s’)\right]$$

Above is the Bellman optimality equation for the state value function.

Linear Programming for MDP

We have a key observation that the Bellman optimality equation can be rewritten as a linear programming problem. The linear programming formulation for MDP is as follows:

$$\begin{array}{rl}\text{minimize}& \sum _{s}V(s)\\ \text{subject to}& V(s)\ge R(s,a)+\gamma \sum _{s’}P(s’|s,a)V(s’),\\ & \forall s\in S,a\in A\end{array}$$

Proof:
The first part is to show that the optimal value function $V^{\ast }$ is a feasible solution to the above linear programming problem. We can see that for any state $s$ and action $a$, we have:$$\begin{gathered} V^{\ast }(s)=\underset{a’}{\max }\left[R(s,a’)+\gamma \sum _{s’}P(s’|s,a’)V^{\ast }(s’)\right] \\ \ge R(s,a)+\gamma \sum _{s’}P(s’|s,a)V^{\ast }(s’) \end{gathered}$$
Thus, $V^{\ast }$ satisfies the constraints of the linear programming problem.
The second part is to show that any feasible solution satisfying the constraints must be greater than or equal to $V^{\ast }$.
Given that the optimal policy ${\pi }^{\ast }$ choose one action ${\pi }^{\ast }(s)$ for each state $s$, LP constraints imply that for any state $s$:$$V(s)\ge R(s,{\pi }^{\ast }(s))+\gamma \sum _{s’}P(s’|s,{\pi }^{\ast }(s))V(s’)$$
Let’s write them in matrix form:$$V\ge R^{{\pi }^{\ast }}+\gamma P^{{\pi }^{\ast }}V$$
while the $R^{{\pi }^{\ast }}$ is the reward vector under policy ${\pi }^{\ast }$, and $P^{{\pi }^{\ast }}$ is the transition matrix under policy ${\pi }^{\ast }$. We can rearrange the above inequality as:$$(I-\gamma P^{{\pi }^{\ast }})V\ge R^{{\pi }^{\ast }}$$
Since $\gamma <1$, we can conclude that $I-\gamma P^{{\pi }^{\ast }}$ is invertible, and we can get its inverse as:$$(I-\gamma P^{{\pi }^{\ast }}{)}^{-1}=\sum _{k=0}^{\infty }(\gamma P^{{\pi }^{\ast }}{)}^k$$
Obviously, the above inverse is a non-negative matrix. Thus, we can multiply both sides of the inequality by $(I-\gamma P^{{\pi }^{\ast }}{)}^{-1}$ to get:$$V\ge (I-\gamma P^{{\pi }^{\ast }}{)}^{-1}{R}^{{\pi }^{\ast }}=V^{{\pi }^{\ast }}$$
Since $V^{\ast }\ge V^{{\pi }^{\ast }}$, we can conclude that $V\ge V^{\ast }$.
So, we have shown that any feasible solution is greater than or equal to $V^{\ast }$, and the optimal value function $V^{\ast }$ is a feasible solution. Thus, the optimal solution to the linear programming problem is $V^{\ast }$.

Given we have the $V^{\ast }$, we can easily derive the optimal policy ${\pi }^{\ast }$ as:$${\pi }^{\ast }(s)=\arg \underset{a}{\max }\left[R(s,a)+\gamma \sum _{s’}P(s’|s,a)V^{\ast }(s’)\right]$$

Here I asked claude to give me a simple explanation of the equivalence between this greedy like policy decision and the optimal policy. The formal proof may be need to use the fixed point theorem, but it is beyond the scope of this post. We just need to remember that the optimal policy can be derived from the optimal value function by choosing the action that maximizes the expected return.

The Dual Linear Programming for MDP

The dual linear programming formulation for MDP is as follows:$$\begin{array}{rl}\text{maximize}& \sum _{s,a}\rho (s,a)R(s,a)\\ \text{subject to}& \sum _a\rho (s’,a)=\sum _{s,a}\rho (s,a)P(s’|s,a),\\ & \forall s’\in S\\ & \rho (s,a)\ge 0,\forall s\in S,a\in A\end{array}$$

If the initial state distribution $\mu (s)>0$ for all $s\in S$, then let
$w(s)=\mu (s)$, the target function means the expected accumulated reward represented by the occupancy measure, and the constraints means the flow conservation constraints.

Assume that the optimal solution is ${\rho }^{\ast }(s,a)$, we can derive the optimal policy from Theorem 2 as:$${\pi }^{\ast }(s)=\frac{{\rho }^{\ast }(s,a)}{{\sum }_a{\rho }^{\ast }(s,a)}$$

Comparison between the Primal and Dual Linear Programming for MDP

Summary

At the beginning of this post, I tried to introduce the MDP-based Reinforcement learning, but I found that the solution of MDP takes a lot of space, so I just introduce the Bellman equation and the linear programming formulation for MDP. In the next post, I will introduce the value iteration and policy iteration algorithms for solving MDP, which are based on the Bellman equation.

The category of decision making problem

Dynamic programming

Dynamic program is used to solve the Sequential decision making problem, feature of this problem is that it’s decision making process is sequential, and the decision at one step will affect the next step, and the reward is received at the end of decision making process, not at each step.

For an example, given a maze like problem below, the agent need to find a way from Position A to Position B, and the time of each way is different. Agent need to find the way with the least time. A simple way to solve this is to list all the possible paths, but if there is a circle, if the map is large, this will be unfeasible.

A better way to solve this is Backward induction, we start from the end point, and for evey point we calculate the time to reach the end point, then we regart the selected point as the new end point, and repeat this process until we reach the start point. This is a dynamic programming method, and it can solve the problem in polynomial time. But due to we need find a backward path, this method is only suitable for DAG, if there is a circle, this method will fail.

The example is just a introduction, we can summarize the features of dynamic programming as follows:

• it start from the end, and caculate the best action for each state.
• it traverse all the states, and for each state, it calculate the best action, and the value of this state.
• it need to define the state, path(state transition), time(online reward)

So it lead to the Principle of Optimality:

An optimal policy has the property that whatever the initial state and initial decision are, the remaining decisions must constitute an optimal policy with regard to the state resulting from the first decision.

Markov Decision Process

Stochastic Process

A stochastic process is a collection of random variables, which can be used to describe the evolution of a system over time. It’s mathematical definition is as follows:$$P(X_{t+1}|X_t,X_{t-1},\dots ,X_0)$$
This means that the probability of the next state $X_{t+1}$ depends on the current state $X_t$ and all the previous states $X_{t-1},\dots ,X_0$.

Markov Process

Compared to stochastic process, Markov process has a stronger assumption, which is “the future is independent of the past given the present”. Mathematically, it’s definition is as follows:$$P(X_{t+1}|X_t,X_{t-1},\dots ,X_0)=P(X_{t+1}|X_t)$$
This means that the probability of the next state $X_{t+1}$ only depends on the current state $X_t$, and is independent of all the previous states $X_{t-1},\dots ,X_0$.

Trying to understand it’s property is that the current state contains all the information about the past, so we can make decision based on the current state without worrying about the past.

Markov Decision Process

Markov Decision Process (MDP) provides a mathematical framework for modeling decision making in situations where the outcome is partly random, partly under the control of a decision maker. An MDP is defined by the following components:

• State space (S): A set of all possible states in the environment.
• Action space (A): A set of all possible actions that the agent can take
• Transition function (P): A function that defines the probability of transitioning from one state to another given a specific action. It is denoted as $P(s’|s,a)$, which represents the probability of transitioning to state $s’$ from state $s$ after taking action $a$.
• Reward function (R): A function that defines the reward received after transitioning from one state to another given a specific action. It is denoted as $R(s,a)$, which represents the reward received after taking action $a$ in state $s$. Sometimes only relates to the State.
• Discount factor (γ): A factor that determines the importance of future rewards. It is a value between 0 and 1, where a value closer to 0 makes the agent prioritize immediate rewards, while a value closer to 1 makes the agent consider future rewards more heavily.

The dynamic feature of MDP

The whole process of MDP is dynamic as follows:

1. The agent observes the current state $s_t$.
2. The agent selects an action $a_t$ based on its policy $\pi (a|s)$, which is a mapping from states to actions.
3. The agent gets a reward $R(s_t,a_t)$.
4. The MDP transitions to a new state $s_{t+1}$ according to the transition function $P(s_{t+1}|s_t,a_t)$.

The total reward that the agent receives over time is often defined as the discounted sum of rewards:$$G_t=R(s_t,a_t)+\gamma R(s_{t+1},a_{t+1})+{\gamma }^{2}R(s_{t+2},a_{t+2})+\dots =\sum _{k=0}^{\infty }{\gamma }^{k}R(s_{t+k},a_{t+k})$$

Markov Policy

In the context of MDP, a policy is a function that depends on the history:$$h_t=(s_0,a_0,s_1,a_1,\dots ,s_{t-1},a_{t-1},s_t)\pi (a_t|h_t)=P(a_t|h_t)$$