minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix hisparse LRU policy (#22170)

Browse Source 
Co-authored-by: huangtingwei9988 <huangtingwei9988@users.noreply.github.com>
Co-authored-by: hzh0425 <hzh0425@users.noreply.github.com>
This commit is contained in:
Zhiqiang Xie
2026-04-05 18:47:58 -07:00
committed by GitHub
co-authored by huangtingwei9988 hzh0425
parent 93109cc89b
commit 41c7c97ff3
1 changed files with 17 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/hisparse.cuh
+17 -14
View File
@@ -281,20 +281,6 @@ __global__ void load_cache_to_device_buffer_kernel(
} }
__syncthreads(); __syncthreads();
// Write back LRU order: evictables at front (LRU), hits at back (MRU).
{
const int total_evictable = HOT_BUFFER_SIZE - s_total_hits;
for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) {
if (i < total_evictable) {
// Evictables: source at backward end, dest at LRU front
req_lru_slots[i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i];
} else {
// Hits: source at forward end, dest at MRU back
req_lru_slots[i] = s_lru_slots_out[i - total_evictable];
}
}
}
// Reset offsets for the miss counting phase (only NUM_TOKEN_CHUNKS + 1 entries needed). // Reset offsets for the miss counting phase (only NUM_TOKEN_CHUNKS + 1 entries needed).
for (int i = tid; i < NUM_TOKEN_CHUNKS + 1; i += BLOCK_SIZE) { for (int i = tid; i < NUM_TOKEN_CHUNKS + 1; i += BLOCK_SIZE) {
s_chunk_offset[i] = 0; s_chunk_offset[i] = 0;
@@ -351,6 +337,23 @@ __global__ void load_cache_to_device_buffer_kernel(
__syncthreads(); __syncthreads();
total_misses = NUM_TOP_K - s_total_hits - s_newest_hit; total_misses = NUM_TOP_K - s_total_hits - s_newest_hit;
// Write back LRU order: evictables at front (LRU), hits at back (MRU).
{
const int total_evictable = HOT_BUFFER_SIZE - s_total_hits;
for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) {
if (i < total_misses) {
// Misses: just loaded from host, place right before hits
req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i];
} else if (i < total_evictable) {
// Remaining evictables: truly stale, dest at LRU front
req_lru_slots[i - total_misses] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i];
} else {
// Hits: source at forward end, dest at MRU back
req_lru_slots[i] = s_lru_slots_out[i - total_evictable];
}
}
}
// each warp copies one miss directly, can be separated into a new kernel if parallelism is a concern // each warp copies one miss directly, can be separated into a new kernel if parallelism is a concern
for (int miss_idx = warp_id; miss_idx < total_misses; miss_idx += NUM_WARPS) { for (int miss_idx = warp_id; miss_idx < total_misses; miss_idx += NUM_WARPS) {
const int32_t miss_token = s_top_k_tokens[miss_idx]; const int32_t miss_token = s_top_k_tokens[miss_idx];