zenlm/zen5-engine
1
0
Fork 0
mirror of https://github.com/zenlm/zen5-engine.git synced 2026-07-26 22:09:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

cuda: use managed KV cache for huge contexts

Browse Source 
Keep normal CUDA context buffers on device allocations, but route very large KV-cache tensors through managed memory so million-token contexts do not starve unified-memory systems during graph/session allocation.

The fallback is scoped to the long-lived KV/cache tensors and logs when it is used because it may reduce performance.

Tested on 0.180 with:
- make cpu
- make -B cuda-spark
- make cuda-regression
- ./ds4_test --server --metal-kernels
- ./ds4_test --logprob-vectors --tool-call-quality
- ds4-bench ctx-alloc 32768, 250000, and 1000000
- ds4-server --ctx 1000000 startup smoke

(cherry picked from commit 0b248a65c07d21f2fc8ff4815bd8b75af26719f9)
This commit is contained in:
antirez
2026-05-14 17:03:51 +02:00
parent 94c1f38314
commit 04b6fda2be
4 changed files with 136 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ds4.c
+75 -4
View File
@@ -8442,6 +8442,47 @@ static bool metal_graph_apply_directional_steering_ffn(
return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f);
}
static uint64_t metal_graph_kv_cache_bytes_for_context(uint32_t ctx_size, uint32_t raw_cap) {
uint64_t bytes = (uint64_t)DS4_N_LAYER *
raw_cap *
DS4_N_HEAD_DIM *
sizeof(float);
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
const uint32_t ratio = ds4_layer_compress_ratio(il);
if (ratio == 0) continue;
const uint64_t comp_cap = (uint64_t)(ctx_size / ratio + 2u);
bytes += comp_cap * DS4_N_HEAD_DIM * sizeof(float);
if (ratio == 4) {
bytes += comp_cap * DS4_N_INDEXER_HEAD_DIM * sizeof(float);
}
}
return bytes;
}
static uint64_t metal_graph_context_bytes_for_kv_policy(
uint32_t ctx_size,
uint32_t raw_cap,
uint32_t prefill_cap,
uint64_t *kv_cache_bytes_out) {
uint32_t min_ratio = UINT32_MAX;
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
const uint32_t ratio = ds4_layer_compress_ratio(il);
if (ratio != 0 && ratio < min_ratio) min_ratio = ratio;
}
if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u;
uint64_t comp_cap = (uint64_t)(ctx_size / min_ratio + 2u);
if (comp_cap < 2u) comp_cap = 2u;
const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context(ctx_size, raw_cap);
if (kv_cache_bytes_out) *kv_cache_bytes_out = kv_cache_bytes;
return kv_cache_bytes + 2ull * comp_cap * prefill_cap * sizeof(float);
}
static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor(bool managed, uint64_t bytes) {
return managed ? ds4_gpu_tensor_alloc_managed(bytes) : ds4_gpu_tensor_alloc(bytes);
}
/* =========================================================================
* Metal Diagnostic Dump Hooks.
* =========================================================================
@@ -8610,6 +8651,28 @@ static bool metal_graph_alloc_raw_cap(
: DS4_N_INDEXER_HEAD_DIM);
const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM;
const uint64_t pc = prefill_cap;
uint64_t kv_cache_bytes = 0;
const uint64_t context_bytes =
metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes);
const bool managed_kv_cache =
ds4_gpu_should_use_managed_kv_cache(kv_cache_bytes, context_bytes) != 0;
if (managed_kv_cache) {
/*
* CUDA device allocations are fastest, but a million-token KV cache is
* large enough to starve DGX Spark's unified CPU/GPU memory once the
* model cache and driver allocations are present. For this one
* long-lived cache class, managed memory restores the old demand-paged
* behavior. It can be slower, but it keeps oversized contexts from
* turning memory pressure into a machine-wide lockup.
*/
fprintf(stderr,
"ds4: CUDA using managed KV cache for ctx=%u "
"(kv cache %.2f GiB, context buffers %.2f GiB); "
"this may degrade performance but is needed for very large contexts\n",
ctx_size,
(double)kv_cache_bytes / 1073741824.0,
(double)context_bytes / 1073741824.0);
}
g->cur_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float));
g->flat_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float));
@@ -8631,13 +8694,17 @@ static bool metal_graph_alloc_raw_cap(
g->kv = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float));
bool state_init_ok = true;
for (uint32_t il = 0; il < DS4_N_LAYER; il++) {
g->layer_raw_cache[il] = ds4_gpu_tensor_alloc((uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float));
g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor(
managed_kv_cache,
(uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float));
const uint32_t ratio = ds4_layer_compress_ratio(il);
if (ratio != 0) {
const uint32_t coff = ratio == 4 ? 2u : 1u;
const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM;
const uint64_t attn_rows = (uint64_t)coff * ratio;
g->layer_attn_comp_cache[il] = ds4_gpu_tensor_alloc((uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * sizeof(float));
g->layer_attn_comp_cache[il] = metal_graph_alloc_kv_cache_tensor(
managed_kv_cache,
(uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * sizeof(float));
g->layer_attn_state_kv[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
g->layer_attn_state_score[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
if (enable_mtp) {
@@ -8658,7 +8725,9 @@ static bool metal_graph_alloc_raw_cap(
if (ratio == 4) {
const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM;
const uint64_t index_rows = (uint64_t)coff * ratio;
g->layer_index_comp_cache[il] = ds4_gpu_tensor_alloc((uint64_t)g->layer_comp_cap[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float));
g->layer_index_comp_cache[il] = metal_graph_alloc_kv_cache_tensor(
managed_kv_cache,
(uint64_t)g->layer_comp_cap[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float));
g->layer_index_state_kv[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
g->layer_index_state_score[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
if (enable_mtp) {
@@ -8727,7 +8796,9 @@ static bool metal_graph_alloc_raw_cap(
g->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float));
g->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float));
g->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float));
g->mtp_raw_cache = ds4_gpu_tensor_alloc((uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float));
g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor(
managed_kv_cache,
(uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float));
g->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float));
g->mtp_n_raw = 0;
}
ds4_cuda.cu
+49
View File
@@ -1302,6 +1302,55 @@ extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
return t;
}
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
if (bytes == 0) bytes = 1;
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
if (!cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes), "managed tensor alloc")) {
free(t);
return NULL;
}
t->bytes = bytes;
t->owner = 1;
return t;
}
static uint64_t cuda_managed_kv_reserve_bytes(uint64_t total_bytes) {
const uint64_t min_reserve = 8ull * 1073741824ull;
const uint64_t max_reserve = 40ull * 1073741824ull;
uint64_t reserve = total_bytes / 4u;
if (reserve < min_reserve) reserve = min_reserve;
if (reserve > max_reserve) reserve = max_reserve;
return reserve;
}
extern "C" int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) {
if (kv_cache_bytes == 0) return 0;
/* Very large KV caches are where device-only cudaMalloc() can make a
* unified-memory machine unresponsive. Managed memory restores the old
* demand-paged behavior for this one long-lived allocation class only. */
const uint64_t huge_kv = 8ull * 1073741824ull;
if (kv_cache_bytes >= huge_kv) return 1;
const uint64_t large_context = 8ull * 1073741824ull;
if (context_bytes < large_context) return 0;
size_t free_b = 0;
size_t total_b = 0;
cudaError_t err = cudaMemGetInfo(&free_b, &total_b);
if (err != cudaSuccess) {
(void)cudaGetLastError();
return 0;
}
const uint64_t free_bytes = (uint64_t)free_b;
const uint64_t total_bytes = (uint64_t)total_b;
const uint64_t reserve_bytes = cuda_managed_kv_reserve_bytes(total_bytes);
if (context_bytes > free_bytes) return 1;
return free_bytes - context_bytes < reserve_bytes;
}
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) {
if (!base || offset > base->bytes || bytes > base->bytes - offset) return NULL;
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
ds4_gpu.h
+2
View File
@@ -19,6 +19,7 @@ int ds4_gpu_init(void);
void ds4_gpu_cleanup(void);
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes);
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes);
ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes);
void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor);
uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor);
@@ -40,6 +41,7 @@ int ds4_gpu_set_model_fd(int fd);
int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size);
int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label);
int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label);
int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes);
void ds4_gpu_set_quality(bool quality);
void ds4_gpu_print_memory_report(const char *label);
ds4_metal.m
+10
View File
@@ -3793,6 +3793,16 @@ ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
}
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
return ds4_gpu_tensor_alloc(bytes);
}
int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) {
(void)kv_cache_bytes;
(void)context_bytes;
return 0;
}
ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) {
if (!base) return NULL;
const DS4MetalTensor *base_obj = ds4_gpu_tensor_const_obj(base);