hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize Qwen3.8 inference on Apple silicon

Browse Source
This commit is contained in:
Georg Bauer
2026-09-04 23:27:09 +02:00
parent bf82df77cb
commit bd6353804b
14 changed files with 10149 additions and 1145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

154
native/metal/ds4_metal.m
View File

@@ -50,6 +50,7 @@ enum {
static id<MTLDevice> g_device;
static id<MTLCommandQueue> g_queue;
static id<MTLLibrary> g_library;
static id<MTLLibrary> g_qwen_library;
static id<MTLCommandBuffer> g_batch_cb;
static id<MTLComputeCommandEncoder> g_batch_enc;
static BOOL g_batch_encoder_concurrent;
@@ -501,6 +502,7 @@ static int g_glm_stream_expert_addr_table_building;
static uint64_t g_model_residency_count;
static int g_model_residency_added_to_queue;
static int g_glm_model_mode;
static int g_qwen_model_mode;
static int g_ssd_streaming_mode;
static int g_glm_streaming_prefill_full_layer_runtime;
static int g_metal4_runtime_available;
@@ -1255,7 +1257,7 @@ static int ds4_gpu_scratch_needs_cpu_access(const char *label) {
static MTLResourceOptions ds4_gpu_model_resource_options(void) {
MTLResourceOptions options = MTLResourceStorageModeShared;
if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
if (g_qwen_model_mode || getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
options |= MTLResourceHazardTrackingModeUntracked;
}
return options;
@@ -2238,7 +2240,10 @@ static id<MTLComputePipelineState> ds4_gpu_get_pipeline(
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name];
id<MTLLibrary> library = strncmp(function_name, "kernel_qwen_", 12) == 0
? g_qwen_library
: g_library;
id<MTLFunction> fn = [library newFunctionWithName:name];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found\n", function_name);
return nil;
@@ -2369,6 +2374,10 @@ static int ds4_gpu_mpp_available(void) {
return g_metal4_tensor_api_enabled && !g_quality_mode;
}
int ds4_gpu_metal4_tensor_api_enabled(void) {
return g_metal4_tensor_api_enabled;
}
/*
* Retained Metal4 defaults live here instead of behind user-visible options.
* The public runtime has one automatic accelerated path plus the global
@@ -4228,6 +4237,10 @@ void ds4_gpu_set_glm_model(bool enabled) {
g_glm_model_mode = enabled ? 1 : 0;
}
void ds4_gpu_set_qwen_model(bool enabled) {
g_qwen_model_mode = enabled ? 1 : 0;
}
void ds4_gpu_set_ssd_streaming(bool enabled) {
g_ssd_streaming_mode = enabled ? 1 : 0;
ds4_gpu_stream_expert_cache_clear_all(1);
@@ -4359,7 +4372,6 @@ static NSString *ds4_gpu_full_source(void) {
@[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"],
@[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"],
@[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.metal"],
@[@"DS4_METAL_QWEN38_SOURCE", @"metal/qwen38.metal"],
@[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"],
@[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"],
@[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"],
@@ -4418,6 +4430,31 @@ static NSString *ds4_gpu_full_source(void) {
return source;
}
static NSString *ds4_gpu_qwen_source(void) {
const char *override_path = getenv("DS4_METAL_QWEN38_SOURCE");
NSMutableArray<NSString *> *paths = [NSMutableArray array];
if (override_path && override_path[0]) {
[paths addObject:[NSString stringWithUTF8String:override_path]];
}
[paths addObject:@"metal/qwen38.metal"];
[paths addObject:@"./metal/qwen38.metal"];
NSFileManager *fm = [NSFileManager defaultManager];
for (NSString *path in paths) {
if (![fm fileExistsAtPath:path]) continue;
NSError *error = nil;
NSString *source = [NSString stringWithContentsOfFile:path
encoding:NSUTF8StringEncoding
error:&error];
if (source) return source;
fprintf(stderr, "ds4: failed to read Metal source %s: %s\n",
[path UTF8String], [[error localizedDescription] UTF8String]);
return nil;
}
fprintf(stderr,
"ds4: Metal source metal/qwen38.metal not found (set DS4_METAL_QWEN38_SOURCE to override)\n");
return nil;
}
typedef struct {
int32_t ne00t;
int32_t ne00;
@@ -6507,6 +6544,37 @@ int ds4_gpu_init(void) {
}
g_library = library;
NSString *qwen_source = ds4_gpu_qwen_source();
if (!qwen_source) {
g_library = nil;
g_queue = nil;
g_device = nil;
return 0;
}
MTLCompileOptions *qwen_options = [MTLCompileOptions new];
if (@available(macOS 15.0, *)) {
qwen_options.mathMode = MTLMathModeSafe;
} else {
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
qwen_options.fastMathEnabled = NO;
#pragma clang diagnostic pop
}
qwen_options.preprocessorMacros = macros;
error = nil;
id<MTLLibrary> qwen_library = [g_device newLibraryWithSource:qwen_source
options:qwen_options
error:&error];
if (!qwen_library) {
fprintf(stderr, "ds4: Qwen Metal shader compilation failed: %s\n",
[[error localizedDescription] UTF8String]);
g_library = nil;
g_queue = nil;
g_device = nil;
return 0;
}
g_qwen_library = qwen_library;
id<MTLFunction> fn = [library newFunctionWithName:@"kernel_get_rows_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n");
@@ -8744,14 +8812,16 @@ void ds4_gpu_test_set_flags(uint32_t flags) {
g_test_flags = flags;
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
static ds4_gpu_tensor *ds4_gpu_tensor_alloc_with_options(
uint64_t bytes,
MTLResourceOptions options) {
if (!g_initialized && !ds4_gpu_init()) return NULL;
if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL;
@autoreleasepool {
DS4MetalTensor *tensor = [DS4MetalTensor new];
tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes
options:MTLResourceStorageModeShared];
options:options];
if (!tensor.buffer) {
return NULL;
}
@@ -8783,6 +8853,16 @@ ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
}
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
return ds4_gpu_tensor_alloc_with_options(bytes, MTLResourceStorageModeShared);
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc_untracked(uint64_t bytes) {
return ds4_gpu_tensor_alloc_with_options(
bytes,
MTLResourceStorageModeShared | MTLResourceHazardTrackingModeUntracked);
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
return ds4_gpu_tensor_alloc(bytes);
}
@@ -10528,6 +10608,7 @@ void ds4_gpu_cleanup(void) {
g_model_buffer_cache = nil;
g_transient_buffers = nil;
g_pending_cbs = nil;
g_qwen_library = nil;
g_library = nil;
g_queue = nil;
g_device = nil;
@@ -11744,13 +11825,14 @@ int ds4_gpu_qwen_dispatch(
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const ds4_gpu_tensor *c,
const ds4_gpu_tensor *d,
const ds4_gpu_qwen_weight_view *weights,
uint32_t weight_count,
const ds4_gpu_qwen_kernel_args *args,
uint32_t grid_x,
uint32_t grid_y) {
if (!kernel || !out || !args || grid_x == 0 || grid_y == 0 ||
weight_count > 3 || (weight_count != 0 && !weights)) {
weight_count > 6 || (weight_count != 0 && !weights)) {
return 0;
}
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
@@ -11775,14 +11857,34 @@ int ds4_gpu_qwen_dispatch(
[enc setBuffer:tensors[i].buffer offset:(NSUInteger)tensors[i].offset atIndex:1 + i];
}
}
if (d) {
const DS4MetalTensor *tensor = ds4_gpu_tensor_const_obj(d);
[enc setBuffer:tensor.buffer offset:(NSUInteger)tensor.offset atIndex:8];
}
for (uint32_t i = 0; i < weight_count; i++) {
uint64_t inner = 0;
id<MTLBuffer> weight = ds4_gpu_wrap_model_exact_range(
weights[i].map,
weights[i].size,
weights[i].offset,
weights[i].bytes,
&inner);
id<MTLBuffer> weight = nil;
if (weights[i].tensor) {
const DS4MetalTensor *tensor = ds4_gpu_tensor_const_obj(weights[i].tensor);
weight = tensor.buffer;
inner = tensor.offset;
} else if (!g_qwen_model_mode &&
ds4_gpu_model_views_cover_range(weights[i].map,
weights[i].size,
weights[i].offset,
weights[i].bytes)) {
weight = ds4_gpu_wrap_model_range(weights[i].map,
weights[i].size,
weights[i].offset,
weights[i].bytes,
&inner);
} else {
weight = ds4_gpu_wrap_model_exact_range(weights[i].map,
weights[i].size,
weights[i].offset,
weights[i].bytes,
&inner);
}
if (!weight) {
ds4_gpu_end_compute_encoder(cb, enc);
if (owned) [cb commit];
@@ -11791,16 +11893,34 @@ int ds4_gpu_qwen_dispatch(
/* Metal resource offsets are four-byte aligned. Preserve an arbitrary
* safetensors data offset for the Qwen kernels to decode explicitly. */
const uint64_t aligned_inner = inner & ~3ull;
bound_args.u[13 + i] = (uint32_t)(inner - aligned_inner);
[enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:5 + i];
static const uint32_t offset_slots[6] = {13, 14, 15, 9, 6, 7};
static const NSUInteger buffer_slots[6] = {5, 6, 7, 9, 10, 11};
const uint32_t offset_slot = offset_slots[i];
const NSUInteger buffer_slot = buffer_slots[i];
bound_args.u[offset_slot] = (uint32_t)(inner - aligned_inner);
[enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:buffer_slot];
}
[enc setBytes:&bound_args length:sizeof(bound_args) atIndex:0];
if (bound_args.u[10] != 0) {
[enc setThreadgroupMemoryLength:bound_args.u[10] atIndex:0];
}
const NSUInteger width = pipeline.threadExecutionWidth;
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
const NSUInteger threads = MIN(MAX(width, 1u), max_threads);
[enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1)
threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
const NSUInteger requested_threads = bound_args.u[12];
if (requested_threads != 0) {
if (requested_threads > max_threads || requested_threads % width != 0) {
ds4_gpu_end_compute_encoder(cb, enc);
if (owned) [cb commit];
return 0;
}
[enc dispatchThreadgroups:MTLSizeMake(grid_x, grid_y, 1)
threadsPerThreadgroup:MTLSizeMake(requested_threads, 1, 1)];
} else {
const NSUInteger threads = MIN(MAX(width, 1u), max_threads);
[enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1)
threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
return owned ? ds4_gpu_finish_command_buffer(cb, 1, kernel) : 1;
}