Load GLM vision weights on demand

This commit is contained in:
Georg Bauer
2026-09-01 20:08:36 +02:00
parent 46d6a976a5
commit ee515ee824
10 changed files with 293 additions and 26 deletions

View File

@@ -1329,7 +1329,7 @@ static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint
}
static id<MTLComputePipelineState> ds4_gpu_get_pipeline(const char *function_name);
static int ds4_gpu_warm_model_views(void);
static int ds4_gpu_warm_model_views(uint32_t first_view);
static double ds4_gpu_gib(uint64_t bytes);
static double ds4_gpu_now_ms(void) {
@@ -1934,7 +1934,8 @@ static int ds4_gpu_add_model_view_range(
static int ds4_gpu_finish_model_views(
double t0,
uint64_t mapped_model_size,
uint64_t display_offset) {
uint64_t display_offset,
uint32_t first_new_view) {
const double t_mapped = ds4_gpu_now_ms();
const int request_residency =
!g_ssd_streaming_mode &&
@@ -1968,7 +1969,7 @@ static int ds4_gpu_finish_model_views(
warmed = 1;
} else {
ds4_gpu_progress_begin("warming Metal model views");
warmed = ds4_gpu_warm_model_views();
warmed = ds4_gpu_warm_model_views(first_new_view);
if (warmed) ds4_gpu_progress_done();
else ds4_gpu_progress_failed();
}
@@ -1994,6 +1995,7 @@ static int ds4_gpu_map_model_views(
uint64_t map_size,
uint64_t max_tensor_bytes) {
const double t0 = ds4_gpu_now_ms();
const uint32_t first_new_view = g_model_view_count;
uint64_t mapped_model_size = 0;
if (!ds4_gpu_add_model_view_range(model_map,
model_size,
@@ -2004,7 +2006,8 @@ static int ds4_gpu_map_model_views(
&mapped_model_size)) {
return 0;
}
return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset);
return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset,
first_new_view);
}
static id<MTLBuffer> ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) {
@@ -2535,8 +2538,8 @@ static void ds4_gpu_detect_metal4_features(void) {
#endif
}
static int ds4_gpu_warm_model_views(void) {
if (g_model_view_count == 0) return 1;
static int ds4_gpu_warm_model_views(uint32_t first_view) {
if (first_view >= g_model_view_count) return 1;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride");
if (!pipeline) return 0;
@@ -2562,7 +2565,7 @@ static int ds4_gpu_warm_model_views(void) {
}
uint64_t total_touches = 0;
for (uint32_t i = 0; i < g_model_view_count; i++) {
for (uint32_t i = first_view; i < g_model_view_count; i++) {
total_touches += (g_model_views[i].bytes + stride - 1) / stride;
}
if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0;
@@ -2585,7 +2588,7 @@ static int ds4_gpu_warm_model_views(void) {
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
uint64_t dst_offset = 0;
for (uint32_t i = 0; i < g_model_view_count; i++) {
for (uint32_t i = first_view; i < g_model_view_count; i++) {
const uint64_t bytes = g_model_views[i].bytes;
const uint64_t n = (bytes + stride - 1) / stride;
[enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0];
@@ -11336,6 +11339,109 @@ int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint
}
}
int ds4_gpu_model_map_active(const void *model_map, uint64_t model_size) {
if (!model_map || model_size == 0) return 0;
for (uint32_t i = 0; i < g_model_view_count; i++) {
if (g_model_views[i].model_map == model_map &&
g_model_views[i].model_size == model_size) return 1;
}
return 0;
}
int ds4_gpu_set_transient_model_map_range(
const void *model_map,
uint64_t model_size,
uint64_t map_offset,
uint64_t map_size,
uint64_t max_tensor_bytes) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!model_map || model_size == 0 || map_offset > model_size || map_size == 0 ||
map_size > model_size - map_offset) return 0;
const uint64_t end = map_offset + map_size;
for (uint32_t i = 0; i < g_model_view_count; i++) {
const uint64_t view_start = g_model_views[i].model_offset;
const uint64_t view_end = view_start + g_model_views[i].bytes;
if (g_model_views[i].model_map == model_map &&
g_model_views[i].model_size == model_size &&
map_offset >= view_start && end <= view_end) return 1;
}
@autoreleasepool {
const uint32_t first_view = g_model_view_count;
uint64_t mapped_size = 0;
max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size,
max_tensor_bytes);
if (ds4_gpu_add_model_view_range(model_map, model_size, map_offset, map_size,
max_tensor_bytes, false, &mapped_size)) {
return 1;
}
while (g_model_view_count > first_view) {
g_model_view_count--;
if (g_model_wrap_count > 0) g_model_wrap_count--;
if (g_model_wrap_bytes >= g_model_views[g_model_view_count].bytes) {
g_model_wrap_bytes -= g_model_views[g_model_view_count].bytes;
} else {
g_model_wrap_bytes = 0;
}
g_model_views[g_model_view_count].buffer = nil;
g_model_views[g_model_view_count].model_map = NULL;
g_model_views[g_model_view_count].model_size = 0;
g_model_views[g_model_view_count].model_offset = 0;
g_model_views[g_model_view_count].bytes = 0;
}
return 0;
}
}
int ds4_gpu_release_transient_model_map(const void *model_map, uint64_t model_size) {
if (!model_map || model_size == 0) return 0;
if (!g_initialized || !ds4_gpu_model_map_active(model_map, model_size)) return 1;
for (uint32_t i = 0; i < g_model_residency_count; i++) {
if (g_model_views[i].model_map == model_map &&
g_model_views[i].model_size == model_size) return 0;
}
@autoreleasepool {
uint32_t kept = 0;
for (uint32_t i = 0; i < g_model_view_count; i++) {
if (g_model_views[i].model_map == model_map &&
g_model_views[i].model_size == model_size) {
if (g_model_wrap_count > 0) g_model_wrap_count--;
if (g_model_wrap_bytes >= g_model_views[i].bytes) {
g_model_wrap_bytes -= g_model_views[i].bytes;
} else {
g_model_wrap_bytes = 0;
}
g_model_views[i].buffer = nil;
g_model_views[i].model_map = NULL;
g_model_views[i].model_size = 0;
g_model_views[i].model_offset = 0;
g_model_views[i].bytes = 0;
continue;
}
if (kept != i) {
g_model_views[kept] = g_model_views[i];
g_model_views[i].buffer = nil;
g_model_views[i].model_map = NULL;
g_model_views[i].model_size = 0;
g_model_views[i].model_offset = 0;
g_model_views[i].bytes = 0;
}
kept++;
}
g_model_view_count = kept;
g_model_wrap_max_bytes = 0;
for (uint32_t i = 0; i < g_model_view_count; i++) {
if (g_model_views[i].bytes > g_model_wrap_max_bytes) {
g_model_wrap_max_bytes = g_model_views[i].bytes;
}
}
return 1;
}
}
static int ds4_gpu_model_views_cover_spans(
const void *model_map,
uint64_t model_size,
@@ -11394,7 +11500,7 @@ int ds4_gpu_set_model_map_spans(
return 0;
}
}
if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) {
if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset, 0)) {
ds4_gpu_model_residency_clear();
ds4_gpu_model_views_clear();
return 0;