diff --git a/opencl/tests/pageatten/cm_pa_common.hpp b/opencl/tests/pageatten/cm_pa_common.hpp index 2872b839..a8fa59a1 100644 --- a/opencl/tests/pageatten/cm_pa_common.hpp +++ b/opencl/tests/pageatten/cm_pa_common.hpp @@ -55,6 +55,8 @@ void pa_lsc_u8( auto q_tokens_left = q_len; static_assert(q_step == REG_N); static_assert(kv_step == REG_K); + static_assert(CMPA_SUB_BLOCK_SZ % 16 == 0); + static_assert(CMPA_BLOCK_SZ % CMPA_SUB_BLOCK_SZ == 0); if (q_tokens_left < 0) q_tokens_left = 0; if (q_tokens_left > q_step) q_tokens_left = q_step; @@ -70,7 +72,10 @@ void pa_lsc_u8( lsc::block_2d_desc b2dK(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); lsc::block_2d_desc b2dV(v_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); - constexpr int quan_blk_stride = CMFLA_NUM_KV_HEADS * (CMFLA_HEAD_SIZE+4) * CMPA_BLOCK_SZ * sizeof(uint8_t); + constexpr int k_quan_blk_stride = CMPA_KVCACHE_U8 == 1 ? + CMFLA_NUM_KV_HEADS * (CMFLA_HEAD_SIZE + 4) * CMPA_BLOCK_SZ * sizeof(uint8_t) : + CMFLA_NUM_KV_HEADS * CMFLA_HEAD_SIZE * (CMPA_BLOCK_SZ + CMPA_BLOCK_SZ / CMPA_SUB_BLOCK_SZ * 4) * sizeof(uint8_t); + constexpr int v_quan_blk_stride = CMFLA_NUM_KV_HEADS * (CMFLA_HEAD_SIZE + 4) * CMPA_BLOCK_SZ * sizeof(uint8_t); int causal_left = q_start+past_lens; constexpr uint slm_buff_size = kv_step * head_size * sizeof(half); @@ -100,8 +105,19 @@ void pa_lsc_u8( } #endif auto cur_block_id = block_indices[kv_pos / CMPA_BLOCK_SZ]; - uint32_t dscale_offset = cur_block_id*quan_blk_stride + \ - CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + kv_pos%CMPA_BLOCK_SZ*sizeof(half); + +#if CMPA_KVCACHE_U8 == 1 + uint32_t k_dscale_offset = cur_block_id * k_quan_blk_stride + CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + + kv_pos % CMPA_BLOCK_SZ * sizeof(half); + uint32_t k_zp_offset = k_dscale_offset + CMPA_BLOCK_SZ * sizeof(half); +#else + uint32_t k_dscale_offset = cur_block_id * k_quan_blk_stride + CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + + (kv_pos % CMPA_BLOCK_SZ) / CMPA_SUB_BLOCK_SZ * head_size *sizeof(half); + uint32_t k_zp_offset = k_dscale_offset + CMPA_BLOCK_SZ / CMPA_SUB_BLOCK_SZ * head_size * sizeof(half); +#endif + uint32_t v_dscale_offset = cur_block_id * v_quan_blk_stride + CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + + kv_pos % CMPA_BLOCK_SZ * sizeof(half); + uint32_t v_zp_offset = v_dscale_offset + CMPA_BLOCK_SZ * sizeof(half); uint slm_offset = (slm_buff_id_write & 3) * slm_buff_size; vector dscale; @@ -110,15 +126,21 @@ void pa_lsc_u8( slm_buff_id_write ++; if (wg_local_id < local_size/2) { - cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset), dscale); - cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset + CMPA_BLOCK_SZ*sizeof(half)), zp); +#if CMPA_KVCACHE_U8 == 1 + cm_svm_block_read(reinterpret_cast( k_cache_base + k_dscale_offset), dscale); + cm_svm_block_read(reinterpret_cast( k_cache_base + k_zp_offset), zp); +#endif matrix kmat; auto quanKmat = kmat.format()[1].format(); - b2dK.set_base_ptr(reinterpret_cast(k_cache_base+cur_block_id*quan_blk_stride)); + b2dK.set_base_ptr(reinterpret_cast(k_cache_base + cur_block_id * k_quan_blk_stride)); b2dK.set_block_y(kv_pos%CMPA_BLOCK_SZ); for(int k = REG_K*wg_local_id; k < head_size; k += REG_K*(local_size/2)) { +#if CMPA_KVCACHE_U8 == 2 + cm_svm_block_read(reinterpret_cast( k_cache_base + k_dscale_offset + k * sizeof(half)), dscale); + cm_svm_block_read(reinterpret_cast( k_cache_base + k_zp_offset + k * sizeof(half)), zp); +#endif cm_load(quanKmat.format(), b2dK.set_block_x(k)); /*@bug: cm compiler in the tail process. : loop combined with type convert. @@ -130,8 +152,13 @@ void pa_lsc_u8( */ #pragma unroll for(int r = 0; r < kv_step; r++) { +#if CMPA_KVCACHE_U8 == 1 kmat[r] = quanKmat[r]-zp[r]; kmat[r] = cm_mul(kmat[r], dscale[r]); +#else + kmat[r] = quanKmat[r]-zp; + kmat[r] = cm_mul(kmat[r], dscale); +#endif } //clear unused data to 0. for(int r = kv_step-1; r >= kv_left; r--) @@ -139,13 +166,13 @@ void pa_lsc_u8( cm_slm_block_write(slm_K, slm_offset + k * kv_step * sizeof(half), kmat.format()); } } else { - cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset), dscale); - cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset+CMPA_BLOCK_SZ*sizeof(half)), zp); + cm_svm_block_read(reinterpret_cast(v_cache_base + v_dscale_offset), dscale); + cm_svm_block_read(reinterpret_cast(v_cache_base + v_zp_offset), zp); matrix VmatVNNI; matrix Vmat; auto quanVmat = Vmat.format().row(1).format(); - b2dV.set_base_ptr(reinterpret_cast(v_cache_base+cur_block_id*quan_blk_stride)); + b2dV.set_base_ptr(reinterpret_cast(v_cache_base + cur_block_id * v_quan_blk_stride)); b2dV.set_block_y(kv_pos%CMPA_BLOCK_SZ); #pragma unroll diff --git a/opencl/tests/pageatten/cm_pa_kernel.hpp b/opencl/tests/pageatten/cm_pa_kernel.hpp index b87a8f81..dec99131 100644 --- a/opencl/tests/pageatten/cm_pa_kernel.hpp +++ b/opencl/tests/pageatten/cm_pa_kernel.hpp @@ -52,6 +52,7 @@ extern "C" _GENX_MAIN_ void cm_page_attention( constexpr int head_size = CMFLA_HEAD_SIZE; constexpr int num_kv_heads = CMFLA_NUM_KV_HEADS; constexpr int pa_block_sz = CMPA_BLOCK_SZ; + constexpr int sub_block_sz = CMPA_SUB_BLOCK_SZ; //# query [q_len, num_heads, S] //# k_cache [kv_len, num_heads, S] //# v_cache [kv_len, num_heads, S] @@ -132,7 +133,12 @@ extern "C" _GENX_MAIN_ void cm_page_attention( #endif #if CMPA_KVCACHE_U8 - uint kv_offset = hkv*(head_size+4)*pa_block_sz; +#if CMPA_KVCACHE_U8 == 1 + uint k_offset = hkv * (head_size + 4) * pa_block_sz; +#else + uint k_offset = hkv * head_size * (pa_block_sz + pa_block_sz / sub_block_sz * 4); +#endif + uint v_offset = hkv * (head_size + 4) * pa_block_sz; pa_lsc_u8( slm_K, slm_V, @@ -143,8 +149,8 @@ extern "C" _GENX_MAIN_ void cm_page_attention( q_len_sg, //q_step, kv_seq_len, //kv_len, reinterpret_cast(query + q_offset), - reinterpret_cast(k_cache + kv_offset), - reinterpret_cast(v_cache + kv_offset), + reinterpret_cast(k_cache + k_offset), + reinterpret_cast(v_cache + v_offset), #if IS_BLOCK_SPARSE reinterpret_cast(block_mask_base), reinterpret_cast(wg_block_mask_base), diff --git a/opencl/tests/pageatten/flashdecode.py b/opencl/tests/pageatten/flashdecode.py new file mode 100644 index 00000000..974db711 --- /dev/null +++ b/opencl/tests/pageatten/flashdecode.py @@ -0,0 +1,487 @@ +from clops.utils import * + +import torch +import torch.nn as nn +import torch.nn.functional as F + +VERBOSE = -1 +enable_vprint = False +def vprint(*all_args): + global enable_vprint + if enable_vprint: + print(*all_args) + +def get_org(Q, K, V, attention_mask): + B,H,L,S = Q.shape + _,Hkv,_,_ = K.shape + out = torch.zeros([B,H,L,S], dtype=Q.dtype) + scale_factor = S**(-0.5) + for b in range(B): + for h in range(H): + hkv = h // (H//Hkv) + attn_score = Q[b, h, :, :].to(dtype=torch.float32) @ (K[b, hkv, :,:].transpose(0,1)).to(dtype=torch.float32) + attn_score *= scale_factor + attn_score += attention_mask[b,0,:,:] + #print(attn_score.shape) + attn_weights = F.softmax(attn_score, 1) + out[b,h,:,:] = attn_weights @ V[b, hkv, :, :].to(dtype=torch.float32) + return out + +# blocking on kv-len dimension with online-softmax +# Softmax(Q@Kt)@V +def get_flash1(query, key, value, attention_mask): + global enable_vprint + B,H,q_len,hs = query.shape + _,Hkv,kv_len,_ = key.shape + out = torch.zeros([B,H,q_len,hs], dtype=value.dtype) + scale_factor = hs**(-0.5) + for b in range(B): + for h in range(H): + hkv = h // (H//Hkv) + Q = query[b, h, :, :] + K = key[b, hkv, :, :] + V = value[b, hkv, :, :] + mask = attention_mask[b,0,:,:] + + # loop one time + # q_len == 1, q_step == 1 + for i in range(0, 1, 1): + i1 = 1 + # online softmax states: + # per-row max-value : [1, 1] + # per-row sum : [1, 1] + # current accumulated V : [1, S] + #cur_max = torch.full([i1-i, 1], torch.finfo(torch.float32).min, dtype=torch.float32) + #cur_sum = torch.full([i1-i, 1], 0, dtype=torch.float32) + #cur_O = torch.full([i1-i, hs], 0, dtype=torch.float32) + + rQ = Q[0, :].reshape(1,hs) # [1,128] sub Q block VNNI packed + + cur_O = torch.full([kv_len // kv_step, 1, hs], 0, dtype=torch.float32) + max_comp_0 = torch.full([kv_len // kv_step], 1, dtype=torch.float32) + + for j in range(0, kv_len, kv_step): + j1 = min(j + kv_step, kv_len) + + if (j == VERBOSE): enable_vprint = True + # compute in local SRAM + # Step8: On chip, compute S(_𝑗)𝑖= Q𝑖K𝑇 π‘—βˆˆ Rπ΅π‘Ÿ ×𝐡𝑐. + rKt = K[j:j1,:].transpose(0,1) #[16, 128] -> [128, 16], suppose kv_step is 16 + rS = (rQ @ rKt).to(dtype=torch.float32).reshape(1,kv_step) # [1,16] + rMask = mask[i:i1, j:j1] # [1,16] + + vprint("rK=", rKt.shape) + vprint("rQt=",rQ.shape) + vprint("rS=",rS.shape) + vprint("rMask=",rMask.shape) + + rS *= scale_factor + rS += rMask + vprint("rS=",rS.shape) + + rowmax = rS.max(1, keepdim=True).values # [1,1] + if j == 0: + cur_max = rowmax + else: + rowmax = torch.maximum(cur_max, rowmax) + vprint("rowmax=", rowmax.shape) + + # compute in local SRAM + rS = torch.exp(rS - rowmax) # [1,16] + vprint("St(Pt)=", rS.shape) + + rowsumP = rS.sum(1, keepdim=True) # [1,1] + vprint("rowsumP=", rowsumP.shape) + + # corrected sum of previous block + if j > 0: + max_comp = torch.exp(cur_max - rowmax) + vprint("max_comp=", max_comp.shape) + max_comp_0[j//kv_step] = max_comp + + if j == 0: + cur_sum = rowsumP + else: + cur_sum = cur_sum * max_comp + rowsumP + + # softmax normalize is saved accoridng to flash-attn2 section 3.1.1 + # We can instead maintain an β€œun-scaled” version of O(2) and keep around the statistics β„“(2) + partial_attn_weight = rS.to(dtype=torch.float16) # [1,16] + + vprint("P=", partial_attn_weight.shape) + + rV = V[j:j1, :] # [16,128] + vprint("rV=",rV.shape) + + # correct last Output to current statistics + cur_O[j//kv_step,:,:] = partial_attn_weight @ rV # [:,1,128] + vprint("cur_O2=", cur_O.shape) + + cur_max = rowmax + if (j == VERBOSE): assert 0 + + cur_O_f32 = cur_O[0,:,:] + for j in range(1, kv_len//kv_step): + cur_O_f32 = cur_O_f32 * max_comp_0[j] + cur_O[j,:,:] + vprint("cur_O_f32=", cur_O_f32.shape) + vprint("cur_sum=", cur_sum.shape) + cur_O_f16 = (cur_O_f32/cur_sum).to(torch.float16) + + if (i == VERBOSE): + enable_vprint = True + print("cur_O_f16=", cur_O_f16.shape, cur_O_f16) + assert 0 + + out[b, h, i:i1, :] = cur_O_f16 + return out + +# Split KV online-softmax +def get_flash2(query, key, value, attention_mask, real_kv_len=0): + global enable_vprint + B,H,q_len,hs = query.shape + _,Hkv,kv_len,_ = key.shape + if real_kv_len == 0: + real_kv_len = kv_len + out = torch.zeros([B,H,q_len,hs], dtype=value.dtype) + scale_factor = hs**(-0.5) + for b in range(B): + for h in range(H): + hkv = h // (H//Hkv) + Q = query[b, h, :, :] + K = key[b, hkv, :, :] + V = value[b, hkv, :, :] + mask = attention_mask[b,0,:,:] + + # loop kv_split + cur_O=torch.full([kv_len//kv_partition_size, 1, hs], 0, dtype=torch.float32) + cur_O_f32=torch.full([1, hs], 0, dtype=torch.float32) + cur_O_f16=torch.full([1, hs], 0, dtype=torch.float16) + lse=torch.full([kv_len//kv_partition_size], 0, dtype=torch.float32) + for i in range(0, kv_len, kv_partition_size): + i1 = min(i + kv_partition_size, kv_len) + # online softmax states: + # per-row max-value : [1, 1] + # per-row sum : [1, 1] + # current accumulated V : [1, S] + #cur_max = torch.full([i1-i, 1], torch.finfo(torch.float32).min, dtype=torch.float32) + #cur_sum = torch.full([i1-i, 1], 0, dtype=torch.float32) + #cur_O = torch.full([i1-i, hs], 0, dtype=torch.float32) + + rQ = Q[0, :].reshape(1,hs) # [1,128] sub Q block VNNI packed + vprint("rQ = ", rQ) + vprint("mask = ", mask) + + cur_lse = 0.0 + cur_sum = 0.0 + for j in range(i, i1, kv_step): + j1 = min(j + kv_step, kv_len) + if j > kv_len: + break + + if (j == VERBOSE): enable_vprint = True + # compute in local SRAM + # Step8: On chip, compute S(_𝑗)𝑖= Q𝑖K𝑇 π‘—βˆˆ Rπ΅π‘Ÿ ×𝐡𝑐. + rKt = K[j:j1,:].transpose(0,1) #[16,128]->[128,16] + rS = (rQ @ rKt).to(dtype=torch.float32).reshape(1,kv_step) #[1,16] + rMask = mask[0:1, j:j1] #[1,16] + + vprint("rK=", rKt) + vprint("rQt=",rQ) + vprint("rS=",rS) + vprint("rMask=",rMask) + + rS *= scale_factor + vprint("rS * scale_factor =",rS) + rS += rMask + + cur_lse += torch.exp(rS).sum(1, keepdim=True).item() # [1,1] + vprint("rS=", rS) + vprint("exp(rS)=", torch.exp(rS)) + vprint("cur_lse=", cur_lse) + + rowmax = rS.max(1, keepdim=True).values # [1,1] + if j == 0: + cur_max = rowmax + else: + rowmax = torch.maximum(cur_max, rowmax) + vprint("rowmax=", rowmax.shape) + + vprint("rS=", rS) + # compute in local SRAM + rS = torch.exp(rS - rowmax) # [1,16] + vprint("rowmax = ", rowmax) + vprint("St(Pt)=", rS) + + rowsumP = rS.sum(1, keepdim=True) # [1,1] + vprint("rowsumP=", rowsumP) + + # corrected sum of previous block + if j > 0: + max_comp = torch.exp(cur_max - rowmax) + vprint("max_comp=", max_comp.shape) + + if j == 0: + cur_sum = rowsumP + else: + cur_sum = cur_sum * max_comp + rowsumP + + # softmax normalize is saved accoridng to flash-attn2 section 3.1.1 + # We can instead maintain an β€œun-scaled” version of O(2) and keep around the statistics β„“(2) + partial_attn_weight = rS.to(dtype=torch.float16) # [1,16] + + vprint("P=", partial_attn_weight) + + rV = V[j:j1, :] # [16,128] + vprint("rV=",rV) + + # correct last Output to current statistics + if j== 0: + cur_O[i//kv_partition_size,:,:] = partial_attn_weight @ rV + else: + cur_O[i//kv_partition_size,:,:] = cur_O[i//kv_partition_size,:,:] * max_comp; + cur_O[i//kv_partition_size,:,:] += partial_attn_weight @ rV # [:,1,128] + vprint("j = ", j) + vprint("cur_O2=", cur_O[i//kv_partition_size,:,:]) + + cur_max = rowmax + if (j == VERBOSE): assert 0 + + lse[i//kv_partition_size] = cur_lse + if i > real_kv_len: + cur_O[i//kv_partition_size,:,:] = 0 + else: + cur_O[i//kv_partition_size,:,:] = cur_O[i//kv_partition_size,:,:] / cur_sum + vprint("cur_sum=", cur_sum.shape) + vprint("cur_O=", cur_O[i//kv_partition_size,:,:]) + + # reduce + # for i in range(0, kv_len//kv_partition_size, 1): + # for j in range(0, hs, kv_step): + # stop = min(j + kv_step, hs) + # print("i=", i, ", j = ", j, ": cur_O[i,:,:]=", cur_O[i,0,j:stop]) + # vprint("lse=", lse) + # print("lse=", lse.shape) # [4] + sum_lse = lse.sum(0) + # print("cur_O=", cur_O.shape) # + # print("cur_O_f32=", cur_O_f32.shape) # + for i in range(0, kv_len//kv_partition_size, 1): + if i * kv_partition_size > real_kv_len: + break + cur_O_f32 += cur_O[i,:,:] * lse[i] / sum_lse + cur_O_f16 = cur_O_f32.to(torch.float16) + out[b, h, :, :] = cur_O_f16 + # print("cur_O_f16=", cur_O_f16.shape) # + # print("out=", out.shape) # + #print("out = ", out[0,0,0,:]) + #print("out = ", out) + return out + +# Split KV online-softmax +def get_flash3(query, key, value, attention_mask): + global enable_vprint + B,H,q_len,hs = query.shape + _,Hkv,kv_len,_ = key.shape + out = torch.zeros([B,H,q_len,hs], dtype=value.dtype) + scale_factor = hs**(-0.5) + for b in range(B): + for h in range(H): + hkv = h // (H//Hkv) + Q = query[b, h, :, :] # [q_len, head_size] + K = key[b, hkv, :, :] # [kv_len, head_size] + V = value[b, hkv, :, :] # [kv_len, head_size] + mask = attention_mask[b,0,:,:] # [q_len, kv_len] + + # loop kv_split + cur_O=torch.full([kv_len//kv_partition_size, 1, hs], 0, dtype=torch.float32) # [kv_len//kv_partition_size, 1, head_size] + cur_O_f32=torch.full([1, hs], 0, dtype=torch.float32) # [1, head_size] + cur_O_f16=torch.full([1, hs], 0, dtype=torch.float16) # [1, head_size] + lse=torch.full([kv_len//kv_partition_size], 0, dtype=torch.float32) # [kv_len//kv_partition_size] + for i in range(0, kv_len, kv_partition_size): + i1 = min(i + kv_partition_size, kv_len) + # online softmax states: + # per-row max-value : [1, 1] + # per-row sum : [1, 1] + # current accumulated V : [1, S] + #cur_max = torch.full([i1-i, 1], torch.finfo(torch.float32).min, dtype=torch.float32) + #cur_sum = torch.full([i1-i, 1], 0, dtype=torch.float32) + #cur_O = torch.full([i1-i, hs], 0, dtype=torch.float32) + + rQ = Q[0, :].reshape(1,hs) # [1,128] sub Q block VNNI packed + if i==0: + vprint("rQ = ", rQ) + vprint("mask = ", mask) + + cur_lse = 0.0 + cur_sum = 0.0 + for j in range(i, i1, kv_partition_size): + j1 = min(j + kv_partition_size, kv_len) + + if (j == VERBOSE): enable_vprint = True + # compute in local SRAM + # Step8: On chip, compute S(_𝑗)𝑖= Q𝑖K𝑇 π‘—βˆˆ Rπ΅π‘Ÿ ×𝐡𝑐. + rKt = K[j:j1,:].transpose(0,1) # [head_size, kv_partition_size] + rS = (rQ @ rKt).to(dtype=torch.float32).reshape(1,kv_partition_size) #[1, kv_partition_size] + rMask = mask[0:1, j:j1] #[1, kv_partition_size] + + vprint("rK=", rKt) + vprint("rQt=",rQ) + vprint("rS=",rS) + vprint("rMask=",rMask) + + rS *= scale_factor + vprint("rS * scale_factor =",rS) + rS += rMask + + cur_lse += torch.exp(rS).sum(1, keepdim=True).item() # [1,1] + vprint("rS=", rS) + vprint("exp(rS)=", torch.exp(rS)) + vprint("cur_lse=", cur_lse) + + rowmax = rS.max(1, keepdim=True).values # [1,1] + + # compute in local SRAM + rS = torch.exp(rS - rowmax) # [1,16] + vprint("rowmax = ", rowmax) + vprint("St(Pt)=", rS) + + rowsumP = rS.sum(1, keepdim=True) # [1,1] + vprint("rowsumP=", rowsumP) + + # corrected sum of previous block + cur_sum = rowsumP + + # softmax normalize is saved accoridng to flash-attn2 section 3.1.1 + # We can instead maintain an β€œun-scaled” version of O(2) and keep around the statistics β„“(2) + partial_attn_weight = rS.to(dtype=torch.float16) # [1,16] + + vprint("P=", partial_attn_weight) + + rV = V[j:j1, :] # [kv_partition_size, head_size] + vprint("rV=",rV) + + # correct last Output to current statistics + cur_O[i//kv_partition_size,:,:] = partial_attn_weight @ rV + vprint("cur_O2=", cur_O[i//kv_partition_size,:,:]) + + lse[i//kv_partition_size] = cur_lse + cur_O[i//kv_partition_size,:,:] = cur_O[i//kv_partition_size,:,:] / cur_sum + vprint("cur_sum=", cur_sum.shape) + vprint("Omat=", cur_O[i//kv_partition_size,:,:]) + + # reduce + # for i in range(0, kv_len//kv_partition_size, 1): + # for j in range(0, hs, kv_step): + # stop = min(j + kv_step, hs) + # print("i=", i, ", j = ", j, ": cur_O[i,:,:]=", cur_O[i,0,j:stop]) + vprint("lse=", torch.log(lse)) + vprint("lse=", lse.shape) # [4] + sum_lse = lse.sum(0) + # print("cur_O=", cur_O.shape) # + # print("cur_O_f32=", cur_O_f32.shape) # + vprint("lse = ", lse) + vprint("sum_lse = ", sum_lse) + for i in range(0, kv_len//kv_partition_size, 1): + cur_O_f32 += cur_O[i,:,:] * lse[i] / sum_lse + cur_O_f16 = cur_O_f32.to(torch.float16) + out[b, h, :, :] = cur_O_f16 + # print("cur_O_f16=", cur_O_f16.shape) # + # print("out=", out.shape) # + #print("out = ", out[0,0,0,:]) + #print("out = ", out) + return out + +if __name__ == "__main__": + import argparse + import os + + torch.manual_seed(0) + torch.set_printoptions(linewidth=1024) + + def check_close(input, other, atol=1e-3, rtol=1e-3): + print(f"[check_close] {input.shape}{input.dtype} vs {other.shape}{other.dtype}") + #print("ref = ", input) + #print("res = ", other) + rtol_max = (((input - other).abs() - 1e-5)/other.abs())[other != 0].max() + atol_max = (((input - other).abs()) - 1e-5*other.abs()).max() + # print(f"[check_close] rtol_max: {rtol_max}") + # print(f"[check_close] atol_max: {atol_max}") + if not torch.allclose(input, other, atol=atol, rtol=rtol): + close_check = torch.isclose(input, other, atol=atol, rtol=rtol) + not_close_indices = torch.where(~close_check) # Invert the close check to find failures + print(f"Not close indices: {not_close_indices}") + print(f" ref_tensor: {input[not_close_indices]}") + print(f" res_tensor: {other[not_close_indices]}") + assert 0 + + parser = argparse.ArgumentParser('') + parser.add_argument('-i', "--impl", type=int, default=1) + parser.add_argument('-b', "--batch", type=int, default=1) + parser.add_argument('-nh', "--num-heads", type=int, default=32) + parser.add_argument('-nkvh', "--num-kv-heads", type=int, default=8) + parser.add_argument('-ql', "--q-len", type=int, default=1) + parser.add_argument('-kvl', "--kv-len", type=int, default=32769) + parser.add_argument('-hs', "--head-size", type=int, default=128) + parser.add_argument('-rkv', "--reset_kv_cache", type=int, default=1) + parser.add_argument('-v', "--verbose", type=int, default=-1) + args = parser.parse_args() + print(args) + + enable_vprint = False + + def vprint(*all_args): + global enable_vprint + if enable_vprint: + print(*all_args) + + batch = args.batch + q_len, q_step = args.q_len, 32 + kv_len, kv_step = args.kv_len, 8 + num_heads = args.num_heads + num_kv_heads = args.num_kv_heads + head_size = args.head_size + enable_gqa = num_heads > num_kv_heads + + # define KV_BLOCK_SIZE = 32,64,128,256 + kv_block_size = 256 + + enable_kvcache_compression = 1 + kv_cache_quantization_mode = os.environ.get("KV_CACHE_QUANT_MODE", "by_token") + kv_cache_quantization_mode = "by_channel" + + def _validate_quant_mode(mode: str) -> str: + mode = mode.strip().lower() + if mode not in {"by_token", "by_channel"}: + raise ValueError(f"Unsupported kv-cache quantization mode: {mode}") + return mode + + kv_cache_quantization_mode = _validate_quant_mode(kv_cache_quantization_mode) + kvcache_quantization_by_token = int(kv_cache_quantization_mode == "by_token") + print(f"{kv_cache_quantization_mode=}, {kvcache_quantization_by_token=}") + + enable_clean_unused_kvcache = args.reset_kv_cache + + low = -127 + high = 128 + act_dtype = torch.float16 + new_kv_len = (kv_len + kv_block_size - 1) // kv_block_size * kv_block_size + q = torch.randint(low, high, [batch, q_len, num_heads, head_size]).to(dtype=act_dtype)/high + k = torch.randint(low, high, [batch, new_kv_len, num_kv_heads, head_size]).to(dtype=act_dtype)/high + v = torch.randint(low, high, [batch, new_kv_len, num_kv_heads, head_size]).to(dtype=act_dtype)/high + attention_mask = torch.full([batch, 1, q_len, kv_len], torch.finfo(act_dtype).min).to(dtype=act_dtype) + + ref = F.scaled_dot_product_attention(q, k[:,:,:kv_len,:], v[:,:,:kv_len,:], attention_mask, dropout_p=0.0, enable_gqa = enable_gqa) + org = get_org(q, k[:,:,:kv_len,:], v[:,:,:kv_len,:], attention_mask) + check_close(ref, org, atol=1e-3, rtol=1e-2) + + org1 = get_flash1(q,k,v,attention_mask) + check_close(ref, org1, atol=1e-3, rtol=1e-2) + print("org of get_flash1 passed !") + + org2 = get_flash2(q,k,v,attention_mask, real_kv_len=kv_len) + check_close(ref, org2, atol=1e-3, rtol=1e-2) + print("org of get_flash2 passed !") + + org = get_flash3(q,k,v,attention_mask) + check_close(ref, org, atol=1e-3, rtol=1e-2) + print("org of get_flash3 passed !") + check_close(org, org2, atol=1e-3, rtol=1e-2) \ No newline at end of file diff --git a/opencl/tests/pageatten/test_pa.py b/opencl/tests/pageatten/test_pa.py index d1a53527..af2077e6 100644 --- a/opencl/tests/pageatten/test_pa.py +++ b/opencl/tests/pageatten/test_pa.py @@ -24,6 +24,39 @@ def get_cm_grf_width(): CM_GRF_WIDTH = get_cm_grf_width() +def quan_per_channel(kv, sub_blk_size=16): + blk_num, kv_heads, blk_size, head_size = kv.shape + + assert blk_size % sub_blk_size == 0, f'Error: blk_size ({blk_size}) must be divisible by sub_blk_size ({sub_blk_size})' + + num_sub_blks = blk_size // sub_blk_size + kv_sub = kv.reshape(blk_num, kv_heads, num_sub_blks, sub_blk_size, head_size) + + # Quantize along the sub_blk_size dimension (dim=3) + # This generates a scale/zp for every channel (head_size) per sub-block + kv_max = kv_sub.amax(dim=3, keepdim=True) + kv_min = kv_sub.amin(dim=3, keepdim=True) + qrange = kv_max - kv_min + + INTMAX = 255.0 + INTMIN = 0.0 + INTRANGE = INTMAX - INTMIN + + # Shape of scale and zp: [blk_num, kv_heads, num_sub_blks, 1, head_size] + kv_scale = (INTRANGE / (qrange + 1e-6)).to(dtype=torch.half) + kv_zp = ((0.0 - kv_min) * kv_scale + INTMIN).to(dtype=torch.half) + + kv_INT8 = torch.round(kv_sub * kv_scale + kv_zp).to(dtype=torch.uint8) + + # Flatten quantized data: [blk_num, kv_heads, blk_size * head_size] + kv_INT8_flat = kv_INT8.reshape(blk_num, kv_heads, -1) + + # Bit-cast metadata to uint8 (each fp16 scale/zp becomes 2 uint8 bytes) + dq_scale_bytes = (1.0 / kv_scale).view(dtype=torch.uint8).reshape(blk_num, kv_heads, -1) + kv_zp_bytes = kv_zp.view(dtype=torch.uint8).reshape(blk_num, kv_heads, -1) + + return torch.concat((kv_INT8_flat, dq_scale_bytes, kv_zp_bytes), dim=-1) + def quan_per_token(kv): blk_num, kv_heads, blksz, *_ = kv.shape kv_max = kv.amax(dim=-1, keepdim = True) @@ -73,7 +106,7 @@ def ALIGN_UP(x, y): def DIV_UP(x, y): return (x + y -1) // y class page_atten_cm: - def __init__(self, num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, is_causal = True, sparse_block_sz = 128): + def __init__(self, num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, sub_block_sz, is_causal = True, sparse_block_sz = 128): self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_size = head_size @@ -84,6 +117,7 @@ def __init__(self, num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compr self.sparse_block_sz = sparse_block_sz self.compressed_kvcache = compressed_kvcache is_block_sparse = True if sparse_block_sz > 1 else False + self.sub_block_sz = sub_block_sz src1 = r'''#include "cm_pa_kernel.hpp"''' cwd = os.path.dirname(os.path.realpath(__file__)) @@ -100,6 +134,7 @@ def __init__(self, num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compr f" -DCMPA_BLOCK_SZ={self.block_sz}" f" -DIS_BLOCK_SPARSE={int(is_block_sparse)}" f" -DCMPA_KVCACHE_U8={int(compressed_kvcache)}" + f" -DCMPA_SUB_BLOCK_SZ={self.sub_block_sz}" f" -mdump_asm -g2") ) @@ -122,7 +157,7 @@ def __call__(self, q, k, v, block_mask, n_repeats = 1): padded_k = torch.nn.functional.pad(k,kv_padding_dims, "constant", 1) padded_v = torch.nn.functional.pad(v,kv_padding_dims, "constant", 1) #padding all to NAN to simulate the NAN case when fp16 - if self.compressed_kvcache == False: + if self.compressed_kvcache == 0: padded_k.view(torch.uint16)[seq_len:aligned_seqlen] = 0xfe00 padded_v.view(torch.uint16)[seq_len:aligned_seqlen] = 0xfe00 @@ -130,15 +165,20 @@ def __call__(self, q, k, v, block_mask, n_repeats = 1): # reorder K,V from [L, H, S] to [block_num, H, block_size, S] k_cache = padded_k.reshape(aligned_seqlen//self.block_sz, self.block_sz, self.num_kv_heads, self.head_size).transpose(1,2).contiguous() v_cache = padded_v.reshape(aligned_seqlen//self.block_sz, self.block_sz, self.num_kv_heads, self.head_size).transpose(1,2).contiguous() - if self.compressed_kvcache: + if self.compressed_kvcache == 1: k_cache = quan_per_token(k_cache) v_cache = quan_per_token(v_cache) + elif self.compressed_kvcache == 2: + k_cache = quan_per_channel(k_cache, self.sub_block_sz) + v_cache = quan_per_token(v_cache) else: k_cache = k_cache.reshape(aligned_seqlen//self.block_sz, self.num_kv_heads, -1) v_cache = v_cache.reshape(aligned_seqlen//self.block_sz, self.num_kv_heads, -1) + #output memory for the whole SDPA output = torch.zeros(seq_len, self.num_heads, self.head_size).to(torch.float16) blks_per_trunk = self.trunk_sz // self.block_sz + assert aligned_seqlen % self.block_sz==0, f'Error: aligned_seqlen must be multiple of block_sz' # Q[L, H, S] # K/V: [blk_num, H, blk_sz, S] @@ -146,8 +186,12 @@ def __call__(self, q, k, v, block_mask, n_repeats = 1): max_blks = aligned_seqlen // self.block_sz kv_dtype = torch.uint8 if self.compressed_kvcache else torch.half - #extra half zp and half scale per token. totally 4 bytes. - token_sz = (head_size+4) if self.compressed_kvcache else (head_size) + # k: per token quantization for compressed_kvcache == 1, per channel quantization for compressed_kvcache == 2 + k_token_sz = (head_size + 4) if self.compressed_kvcache == 1 else (head_size) + k_block_sz = (self.block_sz + self.block_sz // self.sub_block_sz * 4) if self.compressed_kvcache == 2 else (self.block_sz) + # v: always per token quantization for non-zero compressed_kvcache + v_token_sz = (head_size + 4) if self.compressed_kvcache else (head_size) + v_block_sz = self.block_sz if self.sparse_block_sz > 1: block_mask_list = [] @@ -196,8 +240,8 @@ def __call__(self, q, k, v, block_mask, n_repeats = 1): block_indices = torch.randperm(blk_num) # block_indices = torch.arange(blk_num) # print(f'==============={block_indices=}') - sub_k = torch.zeros(blk_num, self.num_kv_heads, self.block_sz*token_sz).to(kv_dtype) - sub_v = torch.zeros(blk_num, self.num_kv_heads, self.block_sz*token_sz).to(kv_dtype) + sub_k = torch.zeros(blk_num, self.num_kv_heads, k_block_sz * k_token_sz).to(kv_dtype) + sub_v = torch.zeros(blk_num, self.num_kv_heads, v_block_sz * v_token_sz).to(kv_dtype) for i in range(len(block_indices)): sub_k[block_indices[i],:] = k_cache[i,:] sub_v[block_indices[i],:] = v_cache[i,:] @@ -266,8 +310,8 @@ def __call__(self, q, k, v, block_mask, n_repeats = 1): @staticmethod @functools.cache - def create_instance(num_heads, num_kv_heads, head_size,block_sz, trunk_sz, compressed_kvcache, is_causal, sparse_block_sz): - return page_atten_cm(num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, is_causal, sparse_block_sz) + def create_instance(num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, sub_block_sz, is_causal, sparse_block_sz): + return page_atten_cm(num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, sub_block_sz, is_causal, sparse_block_sz) # sparse to dense mask def block_mask_to_attention_mask(block_mask: torch.Tensor, q_len: int, kv_len: int, sparse_block_size: int, trunk_sz: int) -> torch.Tensor: @@ -397,7 +441,7 @@ def count_false_percentage(mask): false_percentage = 0.0 return false_percentage -def test_page_attn_causal_batch1(seq_len, num_heads = 16, num_kv_heads = 16, head_size = 80, block_sz=128, trunk_sz=512, compressed_kvcache=False, sparse_block_sz=128, sparse_ratio=0.5, check_acc = True): +def test_page_attn_causal_batch1(seq_len, num_heads = 16, num_kv_heads = 16, head_size = 80, block_sz=128, trunk_sz=512, compressed_kvcache=0, sub_block_sz=16, sparse_block_sz=128, sparse_ratio=0.5, check_acc = True): cl.profiling(True) torch.manual_seed(0) torch.set_printoptions(linewidth=1024) @@ -462,7 +506,7 @@ def generate_block_mask_with_ratio(num_heads, seq_len, trunk_sz, true_ratio=spar density = 1.0 - percentage / 100.0 is_causal = True # PageAttention implictly means causal_mask - pa_cm = page_atten_cm.create_instance(num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, is_causal, sparse_block_sz) + pa_cm = page_atten_cm.create_instance(num_heads, num_kv_heads, head_size, block_sz, trunk_sz, compressed_kvcache, sub_block_sz, is_causal, sparse_block_sz) out = pa_cm(q, k, v, approx_simple_mask) latency = cl.finish() @@ -548,7 +592,7 @@ def check_tril_all(tensor_dir, file_suffix, block_mask_shape): print(f'checked {files_checked} files') return is_tril & files_checked > 0 - compressed_kvcache = True + compressed_kvcache = 1 xattn_thresh = 100 sparse_block_sz, kv_block_size, trunk_sz = 128, 256, 4096 # trunk_sz no use num_heads, num_kv_heads, head_size = 64, 8, 128 @@ -626,7 +670,7 @@ def check_sanity(kv_cache_zps): # print(f'{block_indices=}') is_causal = True - pa_cm = page_atten_cm.create_instance(num_heads, num_kv_heads, head_size, kv_block_size, trunk_sz, compressed_kvcache, is_causal, sparse_block_sz) + pa_cm = page_atten_cm.create_instance(num_heads, num_kv_heads, head_size, kv_block_size, trunk_sz, compressed_kvcache, sub_block_sz, is_causal, sparse_block_sz) t_query = cl.tensor(query.detach().numpy()) t_key_cache = cl.tensor(key_cache.detach().numpy()) @@ -696,30 +740,42 @@ def check_sanity(kv_cache_zps): if __name__ == "__main__": - # test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=True, sparse_block_sz = sparse_block_sz, sparse_ratio=sparse_ratio, check_acc=True) + # test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=1, sparse_block_sz = sparse_block_sz, sparse_ratio=sparse_ratio, check_acc=True) #ACC test PA base if 0: for block_sz in range(32, 144, 16): for blocks_per_trunk in range(1, 30, 6): for seq_len in range(8192, 8248, 3): - for compressed_kv in [False, True]: + for compressed_kv in [0, 1]: print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") print(f'[PA_BASE_ACC_TETS]: seq_len={seq_len} block_sz={block_sz} blocks_per_trunk={blocks_per_trunk} kv_cache=={"U8" if compressed_kv else "F16"} sparse_block_sz=1') print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") - test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kv, sparse_block_sz = 1, check_acc=True) - test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kv, sparse_block_sz = 1, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kv, sub_block_sz=block_sz, sparse_block_sz = 1, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kv, sub_block_sz=block_sz, sparse_block_sz = 1, check_acc=True) for block_sz in range(128, 257, 32): for seq_len in range(32768, 32810): for trunk_num in range(1, 21): - for compressed_kvcache in [True,False,]: + for compressed_kvcache in [1,0,]: seq_in_blks = (seq_len + block_sz -1 ) // block_sz blocks_per_trunk = seq_in_blks // trunk_num if seq_in_blks % trunk_num == 0 else seq_in_blks // (trunk_num - 1) print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") print(f'[PA_BASE_ACC_TETS]:seq_len={seq_len} block_sz={block_sz} blocks_per_trunk={blocks_per_trunk} kv_cache={"U8" if compressed_kvcache else "F16"}') print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") - test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 128, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sparse_block_sz=1, check_acc=True) - test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 128, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sparse_block_sz=1, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 128, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sub_block_sz=block_sz, sparse_block_sz=1, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 128, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sub_block_sz=block_sz, sparse_block_sz=1, check_acc=True) + #ACC test PA per channel + if 1: + seq_len = 32 * 1024 + block_sz = 256 + trunk_sz = seq_len + compressed_kv = 2 + for sub_block_sz in [16, 32, 64, 128]: + print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") + print(f'[PA_BASE_ACC_TETS]: seq_len={seq_len} block_sz={block_sz} trunk_sz={trunk_sz} kv_cache=={"U8" if compressed_kv else "F16"} sparse_block_sz=1') + print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=compressed_kv, sub_block_sz=sub_block_sz, sparse_block_sz = 1, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=compressed_kv, sub_block_sz=sub_block_sz, sparse_block_sz = 1, check_acc=True) #ACC test sparse X Attention: if 0: for sparse_block_sz in [128, 256, 64,]: @@ -727,27 +783,28 @@ def check_sanity(kv_cache_zps): for sparse_ratio in [0.5, 0.75]: for blocks_per_trunk in [1, 15, 16, 17, 32, 300]: for seq_len in [16*15, 16*16, 16*16+1, 1024, 1024+1, 8*1024, 8*1024+3, 16*1024]: - for compressed_kvcache in [True,False,]: + for compressed_kvcache in [1,0,]: print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") print(f'[XATTENION_ACC_TETS]:seq_len={seq_len} block_sz={block_sz} blocks_per_trunk={blocks_per_trunk} kv_cache={"U8" if compressed_kvcache else "F16"} {sparse_block_sz=} {sparse_ratio=}') print("----------------------------------------------------------------------------------------------------------------------------------------------------------------------") - test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sparse_block_sz = sparse_block_sz, sparse_ratio=sparse_ratio, check_acc=True) + test_page_attn_causal_batch1(seq_len, num_heads = 1, num_kv_heads = 1, head_size = 32, block_sz=block_sz, trunk_sz=blocks_per_trunk*block_sz, compressed_kvcache=compressed_kvcache, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=sparse_ratio, check_acc=True) # perf for sparse X attention. - if 0: + if 1: seq_len = 32*1024 block_sz = 256 trunk_sz=seq_len sparse_block_sz = 128 - test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=False, sparse_block_sz = sparse_block_sz, sparse_ratio=0.5, check_acc=False) - # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=True, sparse_block_sz = sparse_block_sz, sparse_ratio=0.5, check_acc=False) + # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=0, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=0.5, check_acc=False) + # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=1, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=0.5, check_acc=False) + test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=2, sub_block_sz=16, sparse_block_sz = sparse_block_sz, sparse_ratio=0.5, check_acc=False) - # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=False, sparse_block_sz = sparse_block_sz, sparse_ratio=0.8, check_acc=False) - # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=True, sparse_block_sz = sparse_block_sz, sparse_ratio=0.8, check_acc=False) + # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=0, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=0.8, check_acc=False) + # test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 4, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=1, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=0.8, check_acc=False) # perf for sparse X attention, with QWen3 8K case - if 1: + if 0: for sparse_block_sz in [128, 256]: for density in [100.0, 0.8, 0.5, 0.25, 0.1]: # seq_len, block_sz, blocks_per_trunk= 8*1024, 256, 16*2 @@ -756,7 +813,7 @@ def check_sanity(kv_cache_zps): print("-----------------------------------------------------------------------------------------------------------------------------------------") print(f'seq_len={seq_len} block_sz={block_sz} blocks_per_trunk={blocks_per_trunk} sparse_block_sz={sparse_block_sz}') print("-----------------------------------------------------------------------------------------------------------------------------------------") - test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=False, sparse_block_sz = sparse_block_sz, sparse_ratio=1.0-density, check_acc=False) - test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=True, sparse_block_sz = sparse_block_sz, sparse_ratio=1.0-density, check_acc=False) + test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=0, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=1.0-density, check_acc=False) + test_page_attn_causal_batch1(seq_len, num_heads = 32, num_kv_heads = 8, head_size = 128, block_sz=block_sz, trunk_sz=trunk_sz, compressed_kvcache=1, sub_block_sz=block_sz, sparse_block_sz = sparse_block_sz, sparse_ratio=1.0-density, check_acc=False) # test_ov()