Skip to content

Per channel quantization for page attention #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xuchen-intel wants to merge 3 commits into
base: dev/quant_by_channel
Choose a base branch
from
Open

Per channel quantization for page attention #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e08682a
add flashdecode to initiate simplify pa_2nd_token.py.
ceciliapeng2011 Jan 15, 2026
0ef10c8
Per channel quantization for page attention
xuchen-intel Jan 12, 2026
294f655
Apply review comments
xuchen-intel Feb 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 36 additions & 9 deletions opencl/tests/pageatten/cm_pa_common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -70,7 +72,10 @@ void pa_lsc_u8(

lsc::block_2d_desc<uint8_t, 1, kv_step, REG_K> b2dK(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0);
lsc::block_2d_desc<uint8_t, 1, REG_K, REG_N> 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);
Expand Down Expand Up @@ -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<half, kv_step> dscale;
Expand All @@ -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<svmptr_t>( k_cache_base + dscale_offset), dscale);
cm_svm_block_read(reinterpret_cast<svmptr_t>( k_cache_base + dscale_offset + CMPA_BLOCK_SZ*sizeof(half)), zp);
#if CMPA_KVCACHE_U8 == 1
cm_svm_block_read(reinterpret_cast<svmptr_t>( k_cache_base + k_dscale_offset), dscale);
cm_svm_block_read(reinterpret_cast<svmptr_t>( k_cache_base + k_zp_offset), zp);
#endif

matrix<half, kv_step, REG_K> kmat;
auto quanKmat = kmat.format<half, 2, kv_step * REG_K/2>()[1].format<uint8_t, kv_step, REG_K>();
b2dK.set_base_ptr(reinterpret_cast<uint8_t*>(k_cache_base+cur_block_id*quan_blk_stride));
b2dK.set_base_ptr(reinterpret_cast<uint8_t*>(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<svmptr_t>( k_cache_base + k_dscale_offset + k * sizeof(half)), dscale);
cm_svm_block_read(reinterpret_cast<svmptr_t>( k_cache_base + k_zp_offset + k * sizeof(half)), zp);
#endif
cm_load<lsc::Normal>(quanKmat.format<uint8_t>(), b2dK.set_block_x(k));
/*@bug: cm compiler in the tail process.
: loop combined with type convert.
Expand All @@ -130,22 +152,27 @@ 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<half>(kmat[r], dscale[r]);
#else
kmat[r] = quanKmat[r]-zp;
kmat[r] = cm_mul<half>(kmat[r], dscale);
#endif
}
//clear unused data to 0.
for(int r = kv_step-1; r >= kv_left; r--)
kmat[r] = 0;
cm_slm_block_write(slm_K, slm_offset + k * kv_step * sizeof(half), kmat.format<half>());
}
} else {
cm_svm_block_read(reinterpret_cast<svmptr_t>(v_cache_base+dscale_offset), dscale);
cm_svm_block_read(reinterpret_cast<svmptr_t>(v_cache_base+dscale_offset+CMPA_BLOCK_SZ*sizeof(half)), zp);
cm_svm_block_read(reinterpret_cast<svmptr_t>(v_cache_base + v_dscale_offset), dscale);
cm_svm_block_read(reinterpret_cast<svmptr_t>(v_cache_base + v_zp_offset), zp);

matrix<half, REG_K/2, REG_N*2> VmatVNNI;
matrix<half, REG_K, REG_N> Vmat;
auto quanVmat = Vmat.format<half, 2, REG_K*REG_N/2>().row(1).format<uint8_t, REG_K, REG_N>();
b2dV.set_base_ptr(reinterpret_cast<uint8_t*>(v_cache_base+cur_block_id*quan_blk_stride));
b2dV.set_base_ptr(reinterpret_cast<uint8_t*>(v_cache_base + cur_block_id * v_quan_blk_stride));
b2dV.set_block_y(kv_pos%CMPA_BLOCK_SZ);

#pragma unroll
Expand Down
12 changes: 9 additions & 3 deletions opencl/tests/pageatten/cm_pa_kernel.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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<is_causal, num_heads, num_kv_heads, head_size, 0>(
slm_K,
slm_V,
Expand All @@ -143,8 +149,8 @@ extern "C" _GENX_MAIN_ void cm_page_attention(
q_len_sg, //q_step,
kv_seq_len, //kv_len,
reinterpret_cast<svmptr_t>(query + q_offset),
reinterpret_cast<svmptr_t>(k_cache + kv_offset),
reinterpret_cast<svmptr_t>(v_cache + kv_offset),
reinterpret_cast<svmptr_t>(k_cache + k_offset),
reinterpret_cast<svmptr_t>(v_cache + v_offset),
#if IS_BLOCK_SPARSE
reinterpret_cast<svmptr_t>(block_mask_base),
reinterpret_cast<svmptr_t>(wg_block_mask_base),
Expand Down
Loading