Per channel quantization for page attention - #16
Conversation
xuchen-intel
changed the title
|
@ceciliapeng2011 Could you please take a look? |
3bc266a to
6b8b600
Compare
4517ccf to
e08682a
Compare
| 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,]: |
There was a problem hiding this comment.
0 for fp16, 1 for per token, 2 for per channel... but why there is no test cases for 2? Have you benchmark them all here to cover any corner case? And is there any test cases to cover different
sub_block_sz?
There was a problem hiding this comment.
Applied. Extra per channel test cases with different sub_block_sz are added. Thanks!
|
|
||
| #if CMPA_KVCACHE_U8 | ||
| uint kv_offset = hkv*(head_size+4)*pa_block_sz; | ||
| uint k_offset = CMPA_KVCACHE_U8 == 1 ? hkv * (head_size + 4) * pa_block_sz : hkv * head_size * (pa_block_sz + pa_block_sz / sub_block_sz * 4); |
There was a problem hiding this comment.
Better use constexpr here.
There was a problem hiding this comment.
and static_assert to check pa_block_sz % sub_block_sz == 0
There was a problem hiding this comment.
Any restriction to the possible values of sub_block_sz ?
There was a problem hiding this comment.
Applied. Compile-time evaluation is now assured. Restrictions about pa_block_sz and sub_block_sz are added.
| 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); | ||
| uint32_t k_dscale_offset = CMPA_KVCACHE_U8 == 1 ? |
There was a problem hiding this comment.
Have you checked assembly here - is there going to be a condition branch inst generated? m.b. constexpr is better?
There was a problem hiding this comment.
No condition branch is found in assembly. But you are right. I've revised it to assure compile-time evaluation. Thanks!
6b8b600 to
294f655
Compare
e08682a to
226cfdb
Compare
2 participants
Implement per channel quantization for page attention.