chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[EAGLE] Fix slow Triton compilation in EAGLE KV cache copy by chunking large num_locs_upper (#15111)

Browse Source
This commit is contained in:
YAMY
2025-12-20 22:14:14 -08:00
committed by GitHub
parent fc05acc2c7
commit 6014365564
1 changed files with 38 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -673,6 +673,9 @@ class MHATokenToKVPool(KVCache):
else:
bytes_per_tile = _KV_COPY_TILE_SIZE_SMALL
# Calculate num_locs_upper to avoid large Triton specialization (e.g. 8192)
chunk_upper = 128 if bytes_per_tile >= _KV_COPY_TILE_SIZE_LARGE else 256
self._kv_copy_config = {
"bytes_per_tile": bytes_per_tile,
"byte_tiles": (stride_bytes + bytes_per_tile - 1) // bytes_per_tile,
@@ -681,9 +684,10 @@ class MHATokenToKVPool(KVCache):
if bytes_per_tile <= _KV_COPY_TILE_SIZE_MEDIUM
else _KV_COPY_NUM_WARPS_LARGE_TILE
),
"num_locs_upper": chunk_upper,
}
dummy_loc = torch.zeros(1, dtype=torch.int32, device=self.device)
dummy_loc = torch.zeros(chunk_upper, dtype=torch.int64, device=self.device)
grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"])
copy_all_layer_kv_cache_tiled[grid](
@@ -692,7 +696,7 @@ class MHATokenToKVPool(KVCache):
dummy_loc,
dummy_loc,
1,
1,
chunk_upper,
BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"],
num_warps=self._kv_copy_config["num_warps"],
num_stages=2,
@@ -902,20 +906,40 @@ class MHATokenToKVPool(KVCache):
), "KV copy not initialized. Set enable_kv_cache_copy=True in __init__"
cfg = self._kv_copy_config
N_upper = next_power_of_2(N)
cap = int(cfg.get("num_locs_upper", 256))
grid = (self.data_ptrs.numel(), cfg["byte_tiles"])
copy_all_layer_kv_cache_tiled[grid](
self.data_ptrs,
self.data_strides,
tgt_loc,
src_loc,
N,
N_upper,
BYTES_PER_TILE=cfg["bytes_per_tile"],
num_warps=cfg["num_warps"],
num_stages=2,
)
if N <= cap:
upper = next_power_of_2(N)
copy_all_layer_kv_cache_tiled[grid](
self.data_ptrs,
self.data_strides,
tgt_loc,
src_loc,
N,
upper,
BYTES_PER_TILE=cfg["bytes_per_tile"],
num_warps=cfg["num_warps"],
num_stages=2,
)
return
# Huge N: chunk, but each chunk's upper is still pow2(<= cap)
for start in range(0, N, cap):
end = min(start + cap, N)
chunk_len = end - start
upper = next_power_of_2(chunk_len)
copy_all_layer_kv_cache_tiled[grid](
self.data_ptrs,
self.data_strides,
tgt_loc[start:end],
src_loc[start:end],
chunk_len,
upper,
BYTES_PER_TILE=cfg["bytes_per_tile"],
num_warps=cfg["num_warps"],
num_stages=2,
)
class MHATokenToKVPoolFP4(MHATokenToKVPool):